[feat] pause-training
jingyaogong
parent
4c3b5bacbf
commit
aebc6fb191
|
|
@ -1,48 +1,39 @@
|
||||||
import argparse
|
|
||||||
import time
|
|
||||||
import math
|
|
||||||
import warnings
|
|
||||||
|
|
||||||
warnings.filterwarnings('ignore')
|
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
import torch
|
|
||||||
import torch.distributed as dist
|
|
||||||
|
|
||||||
__package__ = "trainer"
|
__package__ = "trainer"
|
||||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import time
|
||||||
|
import warnings
|
||||||
|
import torch
|
||||||
|
import torch.distributed as dist
|
||||||
from contextlib import nullcontext
|
from contextlib import nullcontext
|
||||||
from torch import optim, nn
|
from torch import optim, nn
|
||||||
from torch.nn.parallel import DistributedDataParallel
|
from torch.nn.parallel import DistributedDataParallel
|
||||||
from torch.utils.data import DataLoader, DistributedSampler
|
from torch.utils.data import DataLoader, DistributedSampler
|
||||||
from transformers import AutoTokenizer, AutoModel
|
from transformers import AutoTokenizer
|
||||||
from model.model_vlm import MiniMindVLM, VLMConfig
|
from model.model_vlm import MiniMindVLM, VLMConfig
|
||||||
from dataset.lm_dataset import VLMDataset
|
from dataset.lm_dataset import VLMDataset
|
||||||
|
from trainer.trainer_utils import *
|
||||||
|
|
||||||
|
warnings.filterwarnings('ignore')
|
||||||
|
|
||||||
|
|
||||||
def Logger(content):
|
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
|
||||||
if not ddp or dist.get_rank() == 0:
|
|
||||||
print(content)
|
|
||||||
|
|
||||||
|
|
||||||
def get_lr(current_step, total_steps, lr):
|
|
||||||
return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
|
|
||||||
|
|
||||||
|
|
||||||
def train_epoch(epoch, wandb):
|
|
||||||
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
for step, (X, Y, loss_mask, pixel_values) in enumerate(train_loader):
|
for step, (X, Y, loss_mask, pixel_values) in enumerate(loader, start=start_step + 1):
|
||||||
X = X.to(args.device)
|
X = X.to(args.device)
|
||||||
Y = Y.to(args.device)
|
Y = Y.to(args.device)
|
||||||
loss_mask = loss_mask.to(args.device)
|
loss_mask = loss_mask.to(args.device)
|
||||||
pixel_values = pixel_values.to(args.device)
|
pixel_values = pixel_values.to(args.device)
|
||||||
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
|
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
|
||||||
for param_group in optimizer.param_groups:
|
for param_group in optimizer.param_groups:
|
||||||
param_group['lr'] = lr
|
param_group['lr'] = lr
|
||||||
|
|
||||||
with ctx:
|
with autocast_ctx:
|
||||||
res = model(X, pixel_values=pixel_values)
|
res = model(X, pixel_values=pixel_values)
|
||||||
loss = loss_fct(
|
loss = loss_fct(
|
||||||
res.logits.view(-1, res.logits.size(-1)),
|
res.logits.view(-1, res.logits.size(-1)),
|
||||||
|
|
@ -64,27 +55,20 @@ def train_epoch(epoch, wandb):
|
||||||
|
|
||||||
optimizer.zero_grad(set_to_none=True)
|
optimizer.zero_grad(set_to_none=True)
|
||||||
|
|
||||||
if step % args.log_interval == 0:
|
if step % args.log_interval == 0 or step == iters - 1:
|
||||||
spend_time = time.time() - start_time
|
spend_time = time.time() - start_time
|
||||||
Logger(
|
current_loss = loss.item() * args.accumulation_steps
|
||||||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
|
current_lr = optimizer.param_groups[-1]['lr']
|
||||||
epoch + 1,
|
eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60
|
||||||
args.epochs,
|
|
||||||
step,
|
|
||||||
iter_per_epoch,
|
|
||||||
loss.item(),
|
|
||||||
optimizer.param_groups[-1]['lr'],
|
|
||||||
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
|
|
||||||
|
|
||||||
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
|
Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:')
|
||||||
wandb.log({"loss": loss,
|
|
||||||
"lr": optimizer.param_groups[-1]['lr'],
|
|
||||||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
|
|
||||||
|
|
||||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min})
|
||||||
|
|
||||||
|
if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
|
||||||
model.eval()
|
model.eval()
|
||||||
moe_path = '_moe' if model_config.use_moe else ''
|
moe_suffix = '_moe' if vlm_config.use_moe else ''
|
||||||
ckp = f'{args.save_dir}/pretrain_vlm_{model_config.hidden_size}{moe_path}.pth'
|
ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth'
|
||||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||||
state_dict = model.module.state_dict()
|
state_dict = model.module.state_dict()
|
||||||
else:
|
else:
|
||||||
|
|
@ -94,115 +78,95 @@ def train_epoch(epoch, wandb):
|
||||||
}
|
}
|
||||||
clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存
|
clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存
|
||||||
torch.save(clean_state_dict, ckp)
|
torch.save(clean_state_dict, ckp)
|
||||||
|
vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer,
|
||||||
|
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
|
||||||
model.train()
|
model.train()
|
||||||
|
|
||||||
|
|
||||||
def init_model(model_config: VLMConfig):
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained('../model', use_fast=True)
|
|
||||||
moe_path = '_moe' if model_config.use_moe else ''
|
|
||||||
# 加载纯语言模型权重
|
|
||||||
ckp = f'{args.save_dir}/llm_{model_config.hidden_size}{moe_path}.pth'
|
|
||||||
model = MiniMindVLM(model_config, vision_model_path="../model/vision_model/clip-vit-base-patch16")
|
|
||||||
state_dict = torch.load(ckp, map_location=args.device)
|
|
||||||
model.load_state_dict(state_dict, strict=False)
|
|
||||||
|
|
||||||
# 冻结除 vision_proj 外的所有参数
|
|
||||||
for name, param in model.named_parameters():
|
|
||||||
if 'vision_proj' not in name:
|
|
||||||
param.requires_grad = False
|
|
||||||
|
|
||||||
Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
|
||||||
|
|
||||||
_, preprocess = model.vision_encoder, model.processor
|
|
||||||
return model.to(args.device), tokenizer, preprocess
|
|
||||||
|
|
||||||
|
|
||||||
def init_distributed_mode():
|
|
||||||
if not ddp: return
|
|
||||||
global ddp_local_rank, DEVICE
|
|
||||||
|
|
||||||
dist.init_process_group(backend="nccl")
|
|
||||||
ddp_rank = int(os.environ["RANK"])
|
|
||||||
ddp_local_rank = int(os.environ["LOCAL_RANK"])
|
|
||||||
ddp_world_size = int(os.environ["WORLD_SIZE"])
|
|
||||||
DEVICE = f"cuda:{ddp_local_rank}"
|
|
||||||
torch.cuda.set_device(DEVICE)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
|
parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
|
||||||
parser.add_argument("--out_dir", type=str, default="../out")
|
parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
|
||||||
parser.add_argument("--epochs", type=int, default=4)
|
parser.add_argument('--save_weight', default='pretrain_vlm', type=str, help="保存权重的前缀名")
|
||||||
parser.add_argument("--batch_size", type=int, default=16)
|
parser.add_argument("--epochs", type=int, default=4, help="训练轮数")
|
||||||
parser.add_argument("--learning_rate", type=float, default=4e-4)
|
parser.add_argument("--batch_size", type=int, default=16, help="batch size")
|
||||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
parser.add_argument("--learning_rate", type=float, default=4e-4, help="初始学习率")
|
||||||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||||
parser.add_argument("--use_wandb", default=False, action="store_true")
|
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-V")
|
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||||
parser.add_argument("--num_workers", type=int, default=8)
|
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||||
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.jsonl")
|
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||||
parser.add_argument("--images_path", type=str, default="../dataset/pretrain_images")
|
parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
|
||||||
parser.add_argument("--ddp", action="store_true")
|
parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
|
||||||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
|
||||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
|
||||||
parser.add_argument("--warmup_iters", type=int, default=0)
|
parser.add_argument('--max_seq_len', default=640, type=int, help="训练的最大截断长度")
|
||||||
parser.add_argument("--log_interval", type=int, default=100)
|
parser.add_argument('--use_moe', default=False, type=bool, help="是否使用MoE")
|
||||||
parser.add_argument("--save_interval", type=int, default=100)
|
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.jsonl", help="训练数据路径")
|
||||||
parser.add_argument('--local_rank', type=int, default=-1)
|
parser.add_argument("--images_path", type=str, default="../dataset/pretrain_images", help="训练图像路径")
|
||||||
parser.add_argument('--hidden_size', default=512, type=int)
|
parser.add_argument('--from_weight', default='llm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
|
||||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
parser.add_argument('--from_resume', default=0, type=int, help="是否自动检测&续训，0否1是")
|
||||||
parser.add_argument('--max_seq_len', default=640, type=int)
|
parser.add_argument('--freeze_llm', default=True, type=bool, help="是否冻结LLM参数（仅训练vision_proj）")
|
||||||
parser.add_argument('--use_moe', default=False, type=bool)
|
parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
|
||||||
|
parser.add_argument("--wandb_project", type=str, default="MiniMind-V-Pretrain", help="wandb项目名")
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
model_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
# ========== 1. 初始化环境和随机种子 ==========
|
||||||
max_seq_len=args.max_seq_len)
|
local_rank = init_distributed_mode()
|
||||||
max_seq_len = model_config.max_seq_len
|
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||||
args.save_dir = os.path.join(args.out_dir)
|
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||||
|
|
||||||
|
# ========== 2. 配置目录、模型参数、检查ckp ==========
|
||||||
os.makedirs(args.save_dir, exist_ok=True)
|
os.makedirs(args.save_dir, exist_ok=True)
|
||||||
os.makedirs(args.out_dir, exist_ok=True)
|
vlm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||||
tokens_per_iter = args.batch_size * max_seq_len
|
max_seq_len=args.max_seq_len, use_moe=args.use_moe)
|
||||||
torch.manual_seed(1337)
|
ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
|
||||||
|
|
||||||
|
# ========== 3. 设置混合精度 ==========
|
||||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||||
|
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||||
|
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||||
|
|
||||||
args.wandb_run_name = f"MiniMind-V Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
# ========== 4. 配wandb ==========
|
||||||
|
|
||||||
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
|
|
||||||
ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run?
|
|
||||||
ddp_local_rank, DEVICE = 0, "cuda:0"
|
|
||||||
if ddp:
|
|
||||||
init_distributed_mode()
|
|
||||||
args.device = torch.device(DEVICE)
|
|
||||||
|
|
||||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
|
||||||
import swanlab as wandb
|
|
||||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
|
||||||
else:
|
|
||||||
wandb = None
|
wandb = None
|
||||||
|
if args.use_wandb and is_main_process():
|
||||||
|
import swanlab as wandb
|
||||||
|
wandb_id = ckp_data.get('wandb_id') if ckp_data else None
|
||||||
|
resume = 'must' if wandb_id else None
|
||||||
|
wandb_run_name = f"MiniMind-V-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||||
|
wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
|
||||||
|
|
||||||
model, tokenizer, preprocess = init_model(model_config)
|
# ========== 5. 定义模型、数据、优化器 ==========
|
||||||
|
model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight,
|
||||||
|
device=args.device, freeze_llm=args.freeze_llm)
|
||||||
train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess,
|
train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess,
|
||||||
image_special_token=model_config.image_special_token,
|
image_special_token=vlm_config.image_special_token,
|
||||||
max_length=max_seq_len)
|
max_length=vlm_config.max_seq_len)
|
||||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||||
train_loader = DataLoader(
|
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||||
train_ds,
|
|
||||||
batch_size=args.batch_size,
|
|
||||||
pin_memory=True,
|
|
||||||
drop_last=False,
|
|
||||||
shuffle=False,
|
|
||||||
num_workers=args.num_workers,
|
|
||||||
sampler=train_sampler
|
|
||||||
)
|
|
||||||
|
|
||||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
|
|
||||||
optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=args.learning_rate)
|
optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=args.learning_rate)
|
||||||
|
|
||||||
if ddp:
|
# ========== 6. 从ckp恢复状态 ==========
|
||||||
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
start_epoch, start_step = 0, 0
|
||||||
model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
|
if ckp_data:
|
||||||
|
model.load_state_dict(ckp_data['model'], strict=False)
|
||||||
|
optimizer.load_state_dict(ckp_data['optimizer'])
|
||||||
|
scaler.load_state_dict(ckp_data['scaler'])
|
||||||
|
start_epoch = ckp_data['epoch']
|
||||||
|
start_step = ckp_data.get('step', 0)
|
||||||
|
|
||||||
iter_per_epoch = len(train_loader)
|
# ========== 7. DDP包模型 ==========
|
||||||
for epoch in range(args.epochs):
|
if dist.is_initialized():
|
||||||
train_epoch(epoch, wandb)
|
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||||||
|
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||||||
|
|
||||||
|
# ========== 8. 开始训练 ==========
|
||||||
|
for epoch in range(start_epoch, args.epochs):
|
||||||
|
train_sampler and train_sampler.set_epoch(epoch)
|
||||||
|
if epoch == start_epoch and start_step > 0: # 第一个epoch且存在检查点
|
||||||
|
batch_sampler = SkipBatchSampler(train_sampler or range(len(train_ds)), args.batch_size, start_step + 1)
|
||||||
|
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||||
|
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||||
|
train_epoch(epoch, loader, len(loader) + start_step + 1, start_step, wandb)
|
||||||
|
else: # 默认从头开始
|
||||||
|
loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||||
|
train_epoch(epoch, loader, len(loader), 0, wandb)
|
||||||
|
|
|
||||||
|
|
@ -1,48 +1,39 @@
|
||||||
import argparse
|
|
||||||
import time
|
|
||||||
import math
|
|
||||||
import warnings
|
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
import torch
|
|
||||||
import torch.distributed as dist
|
|
||||||
|
|
||||||
__package__ = "trainer"
|
__package__ = "trainer"
|
||||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import time
|
||||||
|
import warnings
|
||||||
|
import torch
|
||||||
|
import torch.distributed as dist
|
||||||
from contextlib import nullcontext
|
from contextlib import nullcontext
|
||||||
from torch import optim, nn
|
from torch import optim, nn
|
||||||
from torch.nn.parallel import DistributedDataParallel
|
from torch.nn.parallel import DistributedDataParallel
|
||||||
from torch.utils.data import DataLoader, DistributedSampler
|
from torch.utils.data import DataLoader, DistributedSampler
|
||||||
from transformers import AutoTokenizer, AutoModel
|
from transformers import AutoTokenizer
|
||||||
from model.model_vlm import MiniMindVLM, VLMConfig
|
from model.model_vlm import MiniMindVLM, VLMConfig
|
||||||
from dataset.lm_dataset import VLMDataset
|
from dataset.lm_dataset import VLMDataset
|
||||||
|
from trainer.trainer_utils import *
|
||||||
|
|
||||||
warnings.filterwarnings('ignore')
|
warnings.filterwarnings('ignore')
|
||||||
|
|
||||||
|
|
||||||
def Logger(content):
|
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
|
||||||
if not ddp or dist.get_rank() == 0:
|
|
||||||
print(content)
|
|
||||||
|
|
||||||
|
|
||||||
def get_lr(current_step, total_steps, lr):
|
|
||||||
return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
|
|
||||||
|
|
||||||
|
|
||||||
def train_epoch(epoch, wandb):
|
|
||||||
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
for step, (X, Y, loss_mask, pixel_values) in enumerate(train_loader):
|
for step, (X, Y, loss_mask, pixel_values) in enumerate(loader, start=start_step + 1):
|
||||||
X = X.to(args.device)
|
X = X.to(args.device)
|
||||||
Y = Y.to(args.device)
|
Y = Y.to(args.device)
|
||||||
loss_mask = loss_mask.to(args.device)
|
loss_mask = loss_mask.to(args.device)
|
||||||
pixel_values = pixel_values.to(args.device)
|
pixel_values = pixel_values.to(args.device)
|
||||||
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
|
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
|
||||||
for param_group in optimizer.param_groups:
|
for param_group in optimizer.param_groups:
|
||||||
param_group['lr'] = lr
|
param_group['lr'] = lr
|
||||||
|
|
||||||
with ctx:
|
with autocast_ctx:
|
||||||
res = model(X, pixel_values=pixel_values)
|
res = model(X, pixel_values=pixel_values)
|
||||||
loss = loss_fct(
|
loss = loss_fct(
|
||||||
res.logits.view(-1, res.logits.size(-1)),
|
res.logits.view(-1, res.logits.size(-1)),
|
||||||
|
|
@ -64,27 +55,20 @@ def train_epoch(epoch, wandb):
|
||||||
|
|
||||||
optimizer.zero_grad(set_to_none=True)
|
optimizer.zero_grad(set_to_none=True)
|
||||||
|
|
||||||
if step % args.log_interval == 0:
|
if step % args.log_interval == 0 or step == iters - 1:
|
||||||
spend_time = time.time() - start_time
|
spend_time = time.time() - start_time
|
||||||
Logger(
|
current_loss = loss.item() * args.accumulation_steps
|
||||||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
|
current_lr = optimizer.param_groups[-1]['lr']
|
||||||
epoch + 1,
|
eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60
|
||||||
args.epochs,
|
|
||||||
step,
|
|
||||||
iter_per_epoch,
|
|
||||||
loss.item(),
|
|
||||||
optimizer.param_groups[-1]['lr'],
|
|
||||||
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
|
|
||||||
|
|
||||||
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
|
Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:')
|
||||||
wandb.log({"loss": loss,
|
|
||||||
"lr": optimizer.param_groups[-1]['lr'],
|
|
||||||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
|
|
||||||
|
|
||||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min})
|
||||||
|
|
||||||
|
if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
|
||||||
model.eval()
|
model.eval()
|
||||||
moe_path = '_moe' if model_config.use_moe else ''
|
moe_suffix = '_moe' if vlm_config.use_moe else ''
|
||||||
ckp = f'{args.save_dir}/sft_vlm_{model_config.hidden_size}{moe_path}.pth'
|
ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth'
|
||||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||||
state_dict = model.module.state_dict()
|
state_dict = model.module.state_dict()
|
||||||
else:
|
else:
|
||||||
|
|
@ -94,116 +78,94 @@ def train_epoch(epoch, wandb):
|
||||||
}
|
}
|
||||||
clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存
|
clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存
|
||||||
torch.save(clean_state_dict, ckp)
|
torch.save(clean_state_dict, ckp)
|
||||||
|
vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer,
|
||||||
|
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
|
||||||
model.train()
|
model.train()
|
||||||
|
|
||||||
|
|
||||||
def init_model(model_config: VLMConfig):
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained('../model')
|
|
||||||
moe_path = '_moe' if model_config.use_moe else ''
|
|
||||||
ckp = f'{args.save_dir}/pretrain_vlm_{model_config.hidden_size}{moe_path}.pth'
|
|
||||||
|
|
||||||
model = MiniMindVLM(model_config, vision_model_path="../model/vision_model/clip-vit-base-patch16")
|
|
||||||
state_dict = torch.load(ckp, map_location=args.device)
|
|
||||||
model.load_state_dict(state_dict, strict=False)
|
|
||||||
|
|
||||||
# 默认全参训练
|
|
||||||
# # 只解冻注意力机制中的投影层参数
|
|
||||||
# for name, param in model.model.named_parameters():
|
|
||||||
# if any(proj in name for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']):
|
|
||||||
# param.requires_grad = True
|
|
||||||
|
|
||||||
Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
|
||||||
|
|
||||||
_, preprocess = model.vision_encoder, model.processor
|
|
||||||
return model.to(args.device), tokenizer, preprocess
|
|
||||||
|
|
||||||
|
|
||||||
def init_distributed_mode():
|
|
||||||
if not ddp: return
|
|
||||||
global ddp_local_rank, DEVICE
|
|
||||||
|
|
||||||
dist.init_process_group(backend="nccl")
|
|
||||||
ddp_rank = int(os.environ["RANK"])
|
|
||||||
ddp_local_rank = int(os.environ["LOCAL_RANK"])
|
|
||||||
ddp_world_size = int(os.environ["WORLD_SIZE"])
|
|
||||||
DEVICE = f"cuda:{ddp_local_rank}"
|
|
||||||
torch.cuda.set_device(DEVICE)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
|
parser = argparse.ArgumentParser(description="MiniMind-V SFT")
|
||||||
parser.add_argument("--out_dir", type=str, default="../out")
|
parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
|
||||||
parser.add_argument("--epochs", type=int, default=2)
|
parser.add_argument('--save_weight', default='sft_vlm', type=str, help="保存权重的前缀名")
|
||||||
parser.add_argument("--batch_size", type=int, default=4)
|
parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
|
||||||
parser.add_argument("--learning_rate", type=float, default=1e-6)
|
parser.add_argument("--batch_size", type=int, default=4, help="batch size")
|
||||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
parser.add_argument("--learning_rate", type=float, default=1e-6, help="初始学习率")
|
||||||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
|
||||||
parser.add_argument("--use_wandb", default=False, action="store_true")
|
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
|
||||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-V")
|
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
|
||||||
parser.add_argument("--num_workers", type=int, default=8)
|
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
|
||||||
parser.add_argument("--data_path", type=str, default="../dataset/sft_data.jsonl")
|
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
|
||||||
parser.add_argument("--images_path", type=str, default="../dataset/sft_images")
|
parser.add_argument("--log_interval", type=int, default=10, help="日志打印间隔")
|
||||||
parser.add_argument("--ddp", action="store_true")
|
parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔")
|
||||||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
|
||||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
|
||||||
parser.add_argument("--warmup_iters", type=int, default=0)
|
parser.add_argument('--max_seq_len', default=1536, type=int, help="训练的最大截断长度")
|
||||||
parser.add_argument("--log_interval", type=int, default=10)
|
parser.add_argument('--use_moe', default=False, type=bool, help="是否使用MoE")
|
||||||
parser.add_argument("--save_interval", type=int, default=10)
|
parser.add_argument("--data_path", type=str, default="../dataset/sft_data.jsonl", help="训练数据路径")
|
||||||
parser.add_argument('--local_rank', type=int, default=-1)
|
parser.add_argument("--images_path", type=str, default="../dataset/sft_images", help="训练图像路径")
|
||||||
parser.add_argument('--hidden_size', default=512, type=int)
|
parser.add_argument('--from_weight', default='pretrain_vlm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
|
||||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
parser.add_argument('--from_resume', default=0, type=int, help="是否自动检测&续训，0否1是")
|
||||||
parser.add_argument('--max_seq_len', default=1536, type=int)
|
parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
|
||||||
parser.add_argument('--use_moe', default=False, type=bool)
|
parser.add_argument("--wandb_project", type=str, default="MiniMind-V-SFT", help="wandb项目名")
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
model_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
# ========== 1. 初始化环境和随机种子 ==========
|
||||||
max_seq_len=args.max_seq_len)
|
local_rank = init_distributed_mode()
|
||||||
max_seq_len = model_config.max_seq_len
|
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
|
||||||
args.save_dir = os.path.join(args.out_dir)
|
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
|
||||||
|
|
||||||
|
# ========== 2. 配置目录、模型参数、检查ckp ==========
|
||||||
os.makedirs(args.save_dir, exist_ok=True)
|
os.makedirs(args.save_dir, exist_ok=True)
|
||||||
os.makedirs(args.out_dir, exist_ok=True)
|
vlm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||||
tokens_per_iter = args.batch_size * max_seq_len
|
max_seq_len=args.max_seq_len, use_moe=args.use_moe)
|
||||||
torch.manual_seed(1337)
|
ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
|
||||||
|
|
||||||
|
# ========== 3. 设置混合精度 ==========
|
||||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||||
|
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||||||
|
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
|
||||||
|
|
||||||
args.wandb_run_name = f"MiniMind-V SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
# ========== 4. 配wandb ==========
|
||||||
|
|
||||||
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
|
|
||||||
ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run?
|
|
||||||
ddp_local_rank, DEVICE = 0, "cuda:0"
|
|
||||||
if ddp:
|
|
||||||
init_distributed_mode()
|
|
||||||
args.device = torch.device(DEVICE)
|
|
||||||
|
|
||||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
|
||||||
import swanlab as wandb
|
|
||||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
|
||||||
else:
|
|
||||||
wandb = None
|
wandb = None
|
||||||
|
if args.use_wandb and is_main_process():
|
||||||
|
import swanlab as wandb
|
||||||
|
wandb_id = ckp_data.get('wandb_id') if ckp_data else None
|
||||||
|
resume = 'must' if wandb_id else None
|
||||||
|
wandb_run_name = f"MiniMind-V-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||||
|
wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
|
||||||
|
|
||||||
model, tokenizer, preprocess = init_model(model_config)
|
# ========== 5. 定义模型、数据、优化器 ==========
|
||||||
|
model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight,
|
||||||
|
device=args.device)
|
||||||
train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess,
|
train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess,
|
||||||
image_special_token=model_config.image_special_token,
|
image_special_token=vlm_config.image_special_token,
|
||||||
max_length=max_seq_len)
|
max_length=vlm_config.max_seq_len)
|
||||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
|
||||||
train_loader = DataLoader(
|
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||||||
train_ds,
|
|
||||||
batch_size=args.batch_size,
|
|
||||||
pin_memory=True,
|
|
||||||
drop_last=False,
|
|
||||||
shuffle=False,
|
|
||||||
num_workers=args.num_workers,
|
|
||||||
sampler=train_sampler
|
|
||||||
)
|
|
||||||
|
|
||||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
|
|
||||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||||
|
|
||||||
if ddp:
|
# ========== 6. 从ckp恢复状态 ==========
|
||||||
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
start_epoch, start_step = 0, 0
|
||||||
model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
|
if ckp_data:
|
||||||
|
model.load_state_dict(ckp_data['model'], strict=False)
|
||||||
|
optimizer.load_state_dict(ckp_data['optimizer'])
|
||||||
|
scaler.load_state_dict(ckp_data['scaler'])
|
||||||
|
start_epoch = ckp_data['epoch']
|
||||||
|
start_step = ckp_data.get('step', 0)
|
||||||
|
|
||||||
iter_per_epoch = len(train_loader)
|
# ========== 7. DDP包模型 ==========
|
||||||
for epoch in range(args.epochs):
|
if dist.is_initialized():
|
||||||
train_epoch(epoch, wandb)
|
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||||||
|
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||||||
|
|
||||||
|
# ========== 8. 开始训练 ==========
|
||||||
|
for epoch in range(start_epoch, args.epochs):
|
||||||
|
train_sampler and train_sampler.set_epoch(epoch)
|
||||||
|
if epoch == start_epoch and start_step > 0: # 第一个epoch且存在检查点
|
||||||
|
batch_sampler = SkipBatchSampler(train_sampler or range(len(train_ds)), args.batch_size, start_step + 1)
|
||||||
|
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||||
|
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
|
||||||
|
train_epoch(epoch, loader, len(loader) + start_step + 1, start_step, wandb)
|
||||||
|
else: # 默认从头开始
|
||||||
|
loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=args.num_workers, pin_memory=True)
|
||||||
|
train_epoch(epoch, loader, len(loader), 0, wandb)
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,158 @@
|
||||||
|
"""
|
||||||
|
训练工具函数集合
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import random
|
||||||
|
import math
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
import torch.distributed as dist
|
||||||
|
from torch.utils.data import Sampler
|
||||||
|
from transformers import AutoTokenizer
|
||||||
|
from model.model_vlm import MiniMindVLM
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def is_main_process():
|
||||||
|
return not dist.is_initialized() or dist.get_rank() == 0
|
||||||
|
|
||||||
|
|
||||||
|
def Logger(content):
|
||||||
|
if is_main_process():
|
||||||
|
print(content)
|
||||||
|
|
||||||
|
|
||||||
|
def get_lr(current_step, total_steps, lr):
|
||||||
|
return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
|
||||||
|
|
||||||
|
|
||||||
|
def init_distributed_mode():
|
||||||
|
if int(os.environ.get("RANK", -1)) == -1:
|
||||||
|
return 0 # 非DDP模式
|
||||||
|
|
||||||
|
dist.init_process_group(backend="nccl")
|
||||||
|
local_rank = int(os.environ["LOCAL_RANK"])
|
||||||
|
torch.cuda.set_device(local_rank)
|
||||||
|
return local_rank
|
||||||
|
|
||||||
|
|
||||||
|
def setup_seed(seed: int):
|
||||||
|
random.seed(seed)
|
||||||
|
np.random.seed(seed)
|
||||||
|
torch.manual_seed(seed)
|
||||||
|
torch.cuda.manual_seed(seed)
|
||||||
|
torch.cuda.manual_seed_all(seed)
|
||||||
|
torch.backends.cudnn.deterministic = True
|
||||||
|
torch.backends.cudnn.benchmark = False
|
||||||
|
|
||||||
|
|
||||||
|
def init_vlm_model(vlm_config, from_weight='pretrain_vlm', tokenizer_path='../model',
|
||||||
|
vision_model_path='../model/vision_model/clip-vit-base-patch16',
|
||||||
|
save_dir='../out', device='cuda', freeze_llm=False):
|
||||||
|
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
||||||
|
model = MiniMindVLM(vlm_config, vision_model_path=vision_model_path)
|
||||||
|
|
||||||
|
if from_weight != 'none':
|
||||||
|
moe_suffix = '_moe' if vlm_config.use_moe else ''
|
||||||
|
weight_path = f'{save_dir}/{from_weight}_{vlm_config.hidden_size}{moe_suffix}.pth'
|
||||||
|
weights = torch.load(weight_path, map_location=device)
|
||||||
|
model.load_state_dict(weights, strict=False)
|
||||||
|
|
||||||
|
# Pretrain阶段：冻结除 vision_proj 外的所有参数
|
||||||
|
if freeze_llm:
|
||||||
|
for name, param in model.named_parameters():
|
||||||
|
if 'vision_proj' not in name:
|
||||||
|
param.requires_grad = False
|
||||||
|
|
||||||
|
# 默认全参训练时的可选配置（已注释）
|
||||||
|
# # 只解冻注意力机制中的投影层参数
|
||||||
|
# for name, param in model.model.named_parameters():
|
||||||
|
# if any(proj in name for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']):
|
||||||
|
# param.requires_grad = True
|
||||||
|
|
||||||
|
Logger(f'所加载VLM Model可训练参数：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||||
|
preprocess = model.processor
|
||||||
|
return model.to(device), tokenizer, preprocess
|
||||||
|
|
||||||
|
|
||||||
|
def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs):
|
||||||
|
os.makedirs(save_dir, exist_ok=True)
|
||||||
|
moe_path = '_moe' if vlm_config.use_moe else ''
|
||||||
|
ckp_path = f'{save_dir}/{weight}_{vlm_config.hidden_size}{moe_path}.pth'
|
||||||
|
resume_path = f'{save_dir}/{weight}_{vlm_config.hidden_size}{moe_path}_resume.pth'
|
||||||
|
|
||||||
|
if model is not None:
|
||||||
|
from torch.nn.parallel import DistributedDataParallel
|
||||||
|
state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict()
|
||||||
|
# 移除vision_encoder参数（不需要保存，因为是预训练的）
|
||||||
|
clean_state_dict = {k: v for k, v in state_dict.items() if not k.startswith('vision_encoder.')}
|
||||||
|
ckp_tmp = ckp_path + '.tmp'
|
||||||
|
torch.save({k: v.half() for k, v in clean_state_dict.items()}, ckp_tmp)
|
||||||
|
os.replace(ckp_tmp, ckp_path)
|
||||||
|
|
||||||
|
wandb_id = None
|
||||||
|
if wandb:
|
||||||
|
if hasattr(wandb, 'get_run'):
|
||||||
|
run = wandb.get_run()
|
||||||
|
wandb_id = getattr(run, 'id', None) if run else None
|
||||||
|
else:
|
||||||
|
wandb_id = getattr(wandb, 'id', None)
|
||||||
|
|
||||||
|
resume_data = {
|
||||||
|
'model': state_dict,
|
||||||
|
'optimizer': optimizer.state_dict(),
|
||||||
|
'epoch': epoch,
|
||||||
|
'step': step,
|
||||||
|
'world_size': dist.get_world_size() if dist.is_initialized() else 1,
|
||||||
|
'wandb_id': wandb_id
|
||||||
|
}
|
||||||
|
for key, value in kwargs.items():
|
||||||
|
if value is not None:
|
||||||
|
if hasattr(value, 'state_dict'):
|
||||||
|
if isinstance(value, DistributedDataParallel):
|
||||||
|
resume_data[key] = value.module.state_dict()
|
||||||
|
else:
|
||||||
|
resume_data[key] = value.state_dict()
|
||||||
|
else:
|
||||||
|
resume_data[key] = value
|
||||||
|
|
||||||
|
resume_tmp = resume_path + '.tmp'
|
||||||
|
torch.save(resume_data, resume_tmp)
|
||||||
|
os.replace(resume_tmp, resume_path)
|
||||||
|
else: # 加载模式
|
||||||
|
if os.path.exists(resume_path):
|
||||||
|
ckp_data = torch.load(resume_path, map_location='cpu')
|
||||||
|
saved_ws = ckp_data.get('world_size', 1)
|
||||||
|
current_ws = dist.get_world_size() if dist.is_initialized() else 1
|
||||||
|
if saved_ws != current_ws:
|
||||||
|
ckp_data['step'] = ckp_data['step'] * saved_ws // current_ws
|
||||||
|
Logger(f'GPU数量变化({saved_ws}→{current_ws})，step已自动转换为{ckp_data["step"]}')
|
||||||
|
return ckp_data
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class SkipBatchSampler(Sampler):
|
||||||
|
def __init__(self, sampler, batch_size, skip_batches=0):
|
||||||
|
self.sampler = sampler
|
||||||
|
self.batch_size = batch_size
|
||||||
|
self.skip_batches = skip_batches
|
||||||
|
|
||||||
|
def __iter__(self):
|
||||||
|
batch = []
|
||||||
|
skipped = 0
|
||||||
|
for idx in self.sampler:
|
||||||
|
batch.append(idx)
|
||||||
|
if len(batch) == self.batch_size:
|
||||||
|
if skipped < self.skip_batches:
|
||||||
|
skipped += 1
|
||||||
|
batch = []
|
||||||
|
continue
|
||||||
|
yield batch
|
||||||
|
batch = []
|
||||||
|
if len(batch) > 0 and skipped >= self.skip_batches:
|
||||||
|
yield batch
|
||||||
|
|
||||||
|
def __len__(self):
|
||||||
|
total_batches = (len(self.sampler) + self.batch_size - 1) // self.batch_size
|
||||||
|
return max(0, total_batches - self.skip_batches)
|
||||||
|
|
||||||
Loading…
Reference in New Issue