我正在使用Omniglot数据集，该数据集是一组19,280张图像，每张图像均为105 x 105（灰度）。
我通过以下转换定义了自定义数据集类：
class OmniglotDataset(Dataset):
def __init__(self, X, transform=None):
self.X = X
self.transform = transform
def __len__(self):
return self.X.shape[0]
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
img = self.X[idx]
if self.transform:
img = self.transform(img)
return img
img_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
X_train.shape
(19280, 105, 105)
train_dataset = OmniglotDataset(X_train, transform=img_transform)
当我索引单个图像时，它返回正确的尺寸：
train_dataset[0].shape
torch.Size([1, 105, 105])
但是，当我索引多个图像时，它会以错误的顺序返回尺寸（我期望是3 x 105 x 105）：
train_dataset[[1,2,3]].shape
torch.Size([105, 3, 105])
您收到错误，是因为尝试将单个图像的转换应用于列表：
获取任何大小的批处理的更方便的方法是使用Dataloader：
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
img_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
omniglot = datasets.Omniglot(root='./data', background=True, download=True, transform = img_transform)
data_loader = DataLoader(omniglot, shuffle=False, batch_size = 8)
for image_batch in data_loader:
# now image_batch contain first eight samples
print(image_batch.shape) # torch.Size([8, 1, 105, 105])
break
如果您确实需要以任意顺序获取图像：
from operator import itemgetter
indexes = [1,3,5]
selected_samples = itemgetter(*b)(omniglot)
本文收集自互联网，转载请注明来源。
如有侵权，请联系[email protected] 删除。
我来说两句