使用 Pytorch 显示每个类的图像数量

Question

我正在将 Pytorch 与 FashionMNIST 数据集一起使用，我想显示来自 10 个类中的每一个的 8 个图像样本。但是，我不知道如何将训练测试拆分为 train_labels，因为我需要在标签（类）上循环并打印每个类的 8 个。 知道如何实现这一目标吗？

classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot')

# Define a transform to normalize the data
transform = transforms.Compose([transforms.ToTensor(),
                              #  transforms.Lambda(lambda x: x.repeat(3,1,1)),
                                transforms.Normalize((0.5, ), (0.5,))])
# Download and load the training data
trainset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data/', download=True, train=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True)
# Download and load the test data
testset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data/', download=True, train=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=True)


print('Training set size:', len(trainset))
print('Test set size:',len(testset))

Answer 1

如果我理解正确，您希望按标签对数据集进行分组，然后显示它们。

您可以先构建一个字典来按标签存储示例：

examples = {i: [] for i in range(len(classes))}

然后遍历训练集并使用标签的索引追加到列表中：

for x, i in trainset:
    examples[i].append(x)

但是，这将涵盖整个系列。如果您想提前停止并避免每个班级收集超过 8 个，您可以通过添加条件来实现：

n_examples = 8
for x, i in trainset:
    if all([len(ex) == n_examples for ex in examples.values()])
        break
    if len(examples[i]) < n_examples:
        examples[i].append(x)

剩下的就是用torchvision.transforms.ToPILImage显示：

transforms.ToPILImage()(examples[3][0])

如果要显示多个，可以使用两个连续的torch.cat，一个在dim=1（按行）上，然后在dim=2（按列）上创建一个网格。

grid = torch.cat([torch.cat(examples[i], dim=1) for i in range(len(classes))], dim=2)
transforms.ToPILImage()(grid)

可能的结果：