在Pytorch上具有相同标签的点的批次

时间:2020-03-17 15:37:08

标签: python classification pytorch

我想在每个包含N个训练点的批次上使用梯度下降训练一个神经网络。我希望这些批次仅包含具有相同标签的点,而不是从训练集中随机取样。

例如,如果我正在使用MNIST进行训练,我希望拥有如下所示的批次:

batch_1 = {0,0,0,0,0,0,0,0}

batch_2 = {3,3,3,3,3,3,3,3}

batch_3 = {7,7,7,7,7,7,7,7}

.....

,依此类推。

我该如何使用pytorch?

1 个答案:

答案 0 :(得分:3)

一种方法是为每个类创建子集和数据加载器,然后通过在每次迭代时在数据加载器之间随机切换来进行迭代:

import torch
from torch.utils.data import DataLoader, Subset
from torchvision.datasets import MNIST
from torchvision import transforms
import numpy as np

dataset = MNIST('path/to/mnist_root/', 
                transform=transforms.ToTensor(),
                download=True)

class_inds = [torch.where(dataset.targets == class_idx)[0]
              for class_idx in dataset.class_to_idx.values()]

dataloaders = [
    DataLoader(
        dataset=Subset(dataset, inds),
        batch_size=8,
        shuffle=True,
        drop_last=False)
    for inds in class_inds]

epochs = 1

for epoch in range(epochs):
    iterators = list(map(iter, dataloaders))   
    while iterators:         
        iterator = np.random.choice(iterators)
        try:
            images, labels = next(iterator)   
            print(labels)
            # do_more_stuff()

        except StopIteration:
            iterators.remove(iterator)

这将适用于任何数据集(不仅限于MNIST)。 这是每次迭代打印标签的结果:

tensor([6, 6, 6, 6, 6, 6, 6, 6])
tensor([3, 3, 3, 3, 3, 3, 3, 3])
tensor([0, 0, 0, 0, 0, 0, 0, 0])
tensor([5, 5, 5, 5, 5, 5, 5, 5])
tensor([8, 8, 8, 8, 8, 8, 8, 8])
tensor([0, 0, 0, 0, 0, 0, 0, 0])
...
tensor([1, 1, 1, 1, 1, 1, 1, 1])
tensor([1, 1, 1, 1, 1, 1])

请注意,通过设置drop_last=False,到处将存在少于batch_size个元素的批次。通过将其设置为True,批次将具有相同的大小,但是将删除一些数据点。