时间戳

首页 > 解决方案 > 在 PyTorch 中定义批量大小 = 1 的手动排序的 MNIST 数据集

pytorch - 在 PyTorch 中定义批量大小 = 1 的手动排序的 MNIST 数据集

问题描述

[] ：这表示一个批次。例如，如果批次大小为 5，那么批次将类似于 [1,4,7,4,2]。[] 的长度表示批量大小。

我想要制作的训练集如下所示：

[1] -> [1] -> [1] -> [1] -> [1] -> [7] -> [7] -> [7] -> [7] -> [7] -> [3] -> [3] -> [3] -> [3] -> [3] -> ... 以此类推

这意味着首先是五个 1（batch size = 1），其次是五个 7s（batch size = 1），第三个是五个 3s（batch size = 1）等等......

有人可以给我一个想法吗？

如果有人可以解释如何用代码实现这一点，那将非常有帮助。

谢谢！:)

标签： pytorchmnistbatchsizepytorch-dataloader

解决方案

如果您想要DataLoader在同一类的一行中返回五个样本，但您不想手动为每个索引定义类，那么您可以创建一个自定义采样器。例如

import torch
import torchvision
import torchvision.transforms as T
from itertools import cycle

class RepeatClassSampler(torch.utils.data.Sampler):
    def __init__(self, targets, repeat_count, length, shuffle=False):
        if not torch.is_tensor(targets):
            targets = torch.tensor(targets)

        self.targets = targets
        self.repeat_count = repeat_count
        self.length = length
        self.shuffle = shuffle

        self.classes = torch.unique(targets).tolist()
        self.class_indices = dict()
        for label in self.classes:
            self.class_indices[label] = torch.nonzero(targets == label).flatten() 

    def __iter__(self):
        class_index_iters = dict()
        for label in self.classes:
            if self.shuffle:
                class_index_iters[label] = cycle(self.class_indices[label][torch.randperm(len(self.class_indices))].tolist())
            else:
                class_index_iters[label] = cycle(self.class_indices[label].tolist())

        if self.shuffle:
            target_iter = cycle(self.targets[torch.randperm(len(self.targets))].tolist())
        else:
            target_iter = cycle(self.targets.tolist())

        def index_generator():
            for i in range(self.length):
                if i % self.repeat_count == 0:
                    current_class = next(target_iter)
                yield next(class_index_iters[current_class])
    
        return index_generator()

    def __len__(self):
        return self.length


mnist = torchvision.datasets.MNIST(root='./', train=True, transform=T.ToTensor())
dataloader = torch.utils.data.DataLoader(
        mnist,
        batch_size=1,
        sampler=RepeatClassSampler(
            targets=mnist.targets,
            repeat_count=5,
            length=15,      # How many total to pick from your dataset
            shuffle=True))

for idx, (x, y) in enumerate(dataloader):
    # training loop
    print(f'Batch {idx+1} labels: {y.tolist()}')

推荐阅读