小编Aym*_*ass的帖子

我正在尝试仅使用 PyTorch Mnist 数据集中的特定数字创建数据加载器

我已经尝试创建自己的采样器，但它不起作用，而且我不确定我是否正确使用了蒙版。

class YourSampler(torch.utils.data.sampler.Sampler):

    def __init__(self, mask):

        self.mask = mask


    def __iter__(self):

        return (self.indices[i] for i in torch.nonzero(self.mask))


    def __len__(self):

        return len(self.mask)


mnist = datasets.MNIST(root=dataroot, train=True, download=True, transform = transform)   

mask = [True if mnist[i][1] == 5 else False for i in range(len(mnist))]

mask = torch.tensor(mask)   

sampler = YourSampler(mask)

trainloader = torch.utils.data.DataLoader(mnist, batch_size=4, sampler = sampler, shuffle=False, num_workers=2)

到目前为止，我遇到了许多不同类型的错误。对于此实现，它是“停止迭代”。我觉得这很简单/愚蠢，但我找不到一个简单的方法来做到这一点。感谢您的帮助！