我正在处理一个图像分类数据集。我的数据集中有31个类别,每个类别都有一个对应的文件夹。训练时,我以如下方式加载数据:
from torchvision import datasets, transformsimport torchdef load_training(root_path, dir, batch_size, kwargs): transform = transforms.Compose( [transforms.Resize([256, 256]), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor()]) data = datasets.ImageFolder(root=root_path + dir, transform=transform) train_loader = torch.utils.data.DataLoader(data, batch_size=batch_size, shuffle=False, drop_last=True, **kwargs) return train_loader
现在对于批量大小为32的情况,每个批次的维度是:[32,3,224,224]
。我知道PyTorch使用PIL来加载图像。但是当我尝试从名为images
的包含32个图像的列表中使用PIL加载图像时,使用以下代码片段:
for i in images: im = Image.open(root_dir + dir + i) im = im.convert('RGB') im.load() im = im.resize((224, 224)) im = np.asarray( im, dtype="int32" ) print(im.shape) temp_im_2 = torch.from_numpy(np.array(temp_im)).float()
它显示im.shape
是(32,224,224,3)
。我哪里做错了?我应该怎么做才能使两种情况下的批次具有相同的形状,比如[32,3,224,224]
?
回答:
为什么不尝试交换轴呢?如果im
有四个维度,可以使用im.transpose(0, 3, 1, 2)
。然而,由于你只加载了一张图像,im.shape
应该返回(224, 224, 3)
,所以im.transpose(2, 0, 1)
应该能将通道维度移到第一位,然后你可以将这些图像堆叠起来形成一个批次。