【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)
posted @ 2020-09-30 16:59  wuliytTaotao  阅读(172)  评论(0编辑  收藏