1,使用背景
当有以下两种情况时:
1) 不知道训练的epoch选取为何值时。过小,训练不充分,泛化能力差; 过大,训练过度,导致过拟合。所以需要动态观察每个epoch后,模型在验证集(也可以不严谨的说是测试集)上的精度,选取精度最大的epoch作为最终的训练结果。
2)在加载的模型基础上继续训练。
在训练模型的时候可能会因为一些问题导致程序中断,或者常常需要观察训练情况的变化来更改学习率等参数,这时候就需要加载中断前保存的模型,并在此基础上继续训练。
2,实战代码
本文以CIFAR10数据集为例,将数据集分为了训练集(5W张)、测试集(1W张),每张图片是3*28*28(CHW)。严格意义上讲,实战项目中还需要有验证集的。这里就将验证集省略掉了,道理都是一样的。
本文将自动检查当前机器是否支持CUDA,自动切换设备
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else