multi-images training 1.0
parent
2b04fce5f6
commit
3e55cd236d
|
|
@ -97,9 +97,9 @@ def train_epoch(epoch, wandb):
|
||||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||||||
model.eval()
|
model.eval()
|
||||||
moe_path = '_moe' if lm_config.use_moe else ''
|
moe_path = '_moe' if lm_config.use_moe else ''
|
||||||
if args.multi:
|
if args.multi: # 多图训练权重保存
|
||||||
ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
|
ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
|
||||||
else:
|
else: # 单图训练权重保存
|
||||||
ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft.pth'
|
ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft.pth'
|
||||||
|
|
||||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue