AcademiCodec
AcademiCodec copied to clipboard
保存权重可以保存成统一的格式
可以例如下面这种格式保存,要不然单机保存的模型根据索引会出现问题,我会在后面提交修复的版本
if epoch % config.common.save_interval == 0:
model_to_save = model.module if config.distributed.data_parallel else model
disc_model_to_save = disc_model.module if config.distributed.data_parallel else disc_model
if not config.distributed.data_parallel or dist.get_rank() == 0:
save_master_checkpoint(epoch, model_to_save, optimizer, scheduler, f'{config.checkpoint.save_location}epoch{epoch}_lr{config.optimization.lr}.pt')
save_master_checkpoint(epoch, disc_model_to_save, optimizer_disc, disc_scheduler, f'{config.checkpoint.save_location}epoch{epoch}_disc_lr{config.optimization.lr}.pt')