如何获得每个 class 的 Pytorch 数据加载器？

Question

我想一次在 1 个 MNIST class 上训练我的模型。

我可以用通用加载器加载数据：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.autograd import Variable


trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
# if not exist, download mnist dataset
root = './data'
train_set = datasets.MNIST(root=root, train=True, transform=trans, download=True)

batch_size = 100

train_loader = torch.utils.data.DataLoader(
                 dataset=train_set,
                 batch_size=batch_size,
                 shuffle=True)

但我不确定如何从这个通用加载器（或最初只有 10 个加载器）创建 10 个加载器（classes/digits 每个加载器 1 个）

Answer 1

一个相当简单的解决方案将涉及按真值对数据集进行分组，并为每个组创建一个唯一的数据加载器：

...
from torch.utils.data import Subset, DataLoader

subsets = {target: Subset(train_set, [i for i, (x, y) in enumerate(train_set) if y == target]) for _, target in train_set.class_to_idx.items()}
loaders = {target: DataLoader(subset) for target, subset in subsets.items()}

然后你可以根据class索引选择一个特定的加载器：

class_3_loader = loaders[3]

如何获得每个 class 的 Pytorch 数据加载器？

How to get a Pytorch data loader per class?

python

pytorch

pytorch-dataloader