pytorch - 如何保存 PyTorch 的 DataLoader 实例？

Question

我想保存 PyTorch 的 torch.utils.data.dataloader.DataLoader例如，这样我就可以从中断的地方继续训练(保持随机种子、状态和所有内容)。

Answer 1

您需要采样器的自定义实现。
可以从以下地址轻松使用:https://gist.github.com/usamec/1b3b4dcbafad2d58faa71a9633eea6a5
您可以像这样保存和恢复:

sampler = ResumableRandomSampler(dataset)
loader = torch.utils.data.DataLoader(dataset, batch_size=2, sampler=sampler, pin_memory=True)

for x in loader:
    print(x)
    break

sampler2 = ResumableRandomSampler(dataset)
torch.save(sampler.get_state(), "test_samp.pth")
sampler2.set_state(torch.load("test_samp.pth"))
loader2 = torch.utils.data.DataLoader(dataset, batch_size=2, sampler=sampler2, pin_memory=True)

for x in loader2:
    print(x)