pytorchでCIFAR10の10クラスデータセットから指定の2クラスだけのデータセットにする方法

やりたいこと

pytorchを使って単純な２クラス分類器の実験をするためにCIFAR10から「犬」「猫」だけのデータセットのdataloaderを生成したいです。

試したこと

cifar10のデータセットをクラス名フォルダに画像として保存して、imagefoldarを使って希望のデータセットを作ることはできました。
が、もう少しシンプルに記述したいです。

コード

コードはほぼチュートリアル通りの記述です。ここから「犬」「猫」のtrainsetかtrainloaderが作れたら理想です。

python
1trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
2                                        download=True, transform=transform)
3trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
4                                          shuffle=True, num_workers=2)

行動規範の内容に同意します

回答1件

自己解決

simplerってpytorchのメソッドを使うことで自己解決できました。

python
1k0 = 2 # for cat
2k1 = 4 # for dog
3mask = [s[1]==k0 or s[1]==k1 for s in trainset.samples]
4sampler = torch.utils.data.sampler.WeightedRandomSampler(mask, len(trainset))
5trainloader0 = torch.utils.data.DataLoader(trainset, batch_size=4,
6                                           sampler=sampler, num_workers=2)
7