使用pytorch数据加载器转储图像数据并加载

Question

我想转储数据，以便可以将其重新加载以训练模型。

我的代码因转储数据而被剪断：

for batch_idx, (image, label) in enumerate(dataloader):
    image, label = image.to(device), label.to(device)
    perturbed_image = attack.perturb(image, label)
    
    #---------- Classifier ----------
    predict_A = classifier(perturbed_image)
    pred_label = torch.max(predict_A.data, 1)[1]
    
    if pred_label != label:
        adv_data.append( (perturbed_image.to("cpu"), label.to("cpu")) )

还有其他方法可以正确地将其转储以加载到torch.utils.data.DataLoader中。

Answer 1

最直接的方法是使用torch.save将perturbed_image和label的实际张量保存为二进制文件，然后使用custom Dataset。请注意，保存的张量不是单个图像/标签，而是成批的图像/标签。您的新自定义Dataset应该可以解决这个问题。