【PyTorch】按照 steps 训练和保存模型

在模型训练过程中，一个 epoch 指遍历一遍训练集，而一般的模型训练也是指定多少个 epoch，每个 epoch 结束后看看模型在验证集上的效果并保存模型。

但在有些场景下，如半监督学习，有标记的样本很少，一个 epoch 甚至只有一个 batch 的数据，这个时候频繁查看验证集效果很耗时。

当数据集很小时，训练多久用 epoch 表示不太合适，这个时候使用模型更新次数来表示更加合理，每多少个 steps 查看一次验证集效果并保存模型。

我们可以通过给 DataLoader 传入一个重复采样的随机采样器 RandomSampler 来实现这个功能，其它代码和按照 epoch 训练一致。

复制# batch_size = 64
# steps_to_save = 1024，每 1024 个 steps 查看验证集效果并保存模型，相当于一个 epoch 有 1024 个 steps，只是数据有重复罢了。
trainloader = DataLoader(
    dataset,
    sampler=torch.utils.data.RandomSampler(
        dataset,
        replacement=True,
        num_samples=64*steps_to_save),
    batch_size=64,
    num_workers=4)