我们在训练模型的时候经常出现各种问题导致训练中断,比方说断电,或者关机之类的导致电脑系统关闭,从而将模型训练中断,那么如何在模型中断后,能够保留之前的训练结果不被丢失,同时又可以继续之前的断点处继续训练?

首先在代码离需要保存模型,比方说我们模型设置训练5000轮,那么我们可以选择每100轮保存一次模型,这样的话,在训练的过程中就能保存下100,200,300.。。。等轮数时候的模型,那么当模型训练到400轮的时候突然训练中断,那么我们就可以通过加载400轮的参数来进行继续训练,其实这个过程就类似在预训练模型的基础上进行训练。下面简单粗暴上代码:

1、保存模型

torch.save(checkpoint, checkpoint_path)

其中checkpoint其实保存的就是模型的一些参数,比方说下面这种字典形式的保存所需的模型参数:

checkpoint = {
    'model': model_state_dict,
    'generator': generator_state_dict,
    'opt': model_opt,
    'optim': optim,
}

checkpoint_path则是表示保存的模型

checkpoint_path = '%s_step_%d.pt' % (self.base_path, step)

save_checkpoint_steps是保存的间隔轮数,step是保存的轮数,比方说save_checkpoint_steps=100,那么step的取值就是100,200,300,400等,下面的代码解释step的取值由来。

if step % self.save_checkpoint_steps != 0:
    return
chkpt, chkpt_name = self._save(step)

其中_save函数就是实现了前面checkpoint的内容的保存。

模型的保存设置就此结束。

2、模型的加载

假如此时模型训练中断了,我们得在代码里设置一个参数,这个参数用来查找确定当前路径下是否有已存在得模型。

# 如果有保存的模型,则加载模型,并在其基础上继续训练
    if os.path.exists(log_dir):
        checkpoint = torch.load(log_dir)
        model.load_state_dict(checkpoint['model'])
        generator.load_state_dict(checkpoint['generator'])
        start_epoch = checkpoint['model_opt']
        optim=checkpoint['optim']
        print('加载 epoch {} 成功!'.format(start_epoch))
    else:
        start_epoch = 0
        print('无保存模型,将从头开始训练!')

或者设置一个变量train_from,若赋值已有模型得路径,则继续训练;若为None,那么从头训练。这块代码既可以用于训练中断,又可以用于使用预训练模型。

if opt.train_from:#是否存在预训练模型
    logger.info('Loading checkpoint from %s' % opt.train_from)
    checkpoint = torch.load(opt.train_from)#加载预训练模型的检查点
    model_opt = checkpoint['opt']
else:
    checkpoint = None
    model_opt = opt

加油,come on!

Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐