保存数据:
atk.save(data_loader, save_path="./data/sample.pt", verbose=True)
加载数据:
import torch
from torch.utils.data import DataLoader, TensorDataset
adv_images, labels = torch.load("./data/sample.pt")
图片显示可以直接输出:(需要转换格式)
def imshow(tensor):
image = tensor.cpu().clone()
image = image.squeeze(0) # [3, H, W]
image = transforms.ToPILImage()(image)
# show an image whose values are between [0, 1]
plt.imshow(image)
plt.show()
def imshow_batch(input, title):
# torch.Tensor => numpy
input = input.numpy().transpose((1, 2, 0))
# display images
plt.imshow(input)
plt.title(title)
plt.show()
可以再输出个标签:
# cifar10
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
out = torchvision.utils.make_grid(adv_images[:4])
imshow_batch(out, title='original labels:' + str([class_names[x] for x in labels[:4]]))
但是,用这种方式将加载的数据输入训练好的模型进行测试时,系统一直在自动重启,原因可能是没有设置batch,因此数据量太大导致内存溢出,所以这种方式并不好(当然你也可以手动设置batch)。
因此,我们研究以dataloader的加载dataset的方式处理我们存储的数据,思路请参考本文的参考,具体实现如下:
adv_images, labels = torch.load("./data/sample.pt")
# 定义GetLoader类,继承Dataset方法,并重写__getitem__()和__len__()方法
class GetLoader(torch.utils.data.Dataset):
# 初始化函数,得到数据
def __init__(self, data_img, data_label):
self.data = data_img
self.label = data_label
# index是根据batchsize划分数据后得到的索引,最后将data和对应的labels进行一起返回
def __getitem__(self, index):
data = self.data[index]
labels = self.label[index]
return data, labels
# 该函数返回数据大小长度,目的是DataLoader方便划分,如果不知道大小,DataLoader会一脸懵逼
def __len__(self):
return len(self.data)
# 通过GetLoader将数据进行加载,返回Dataset对象,包含data和labels
PGDdata = GetLoader(adv_images, labels)
#dataloader
adv_dataloader = DataLoader(PGDdata,batch_size =64, shuffle=True,num_workers=4)
然后我们就可以正常进行操作啦~~
参考:Torch load载入数据集的标准流程_Hobbitlzy的博客-CSDN博客_torch 加载数据集 Pytorch加载自己的数据集(使用DataLoader读取Dataset)_l8947943s的博客-CSDN博客