PyTorchによるMNISTデータセット(0,6の2種類の数字だけ)の用意について

Question

前提・実現したいこと

こんにちは．私は現在PyTorchによるMNIST手書き数字のデータセットを用いた深層学習に取り組んでいます．
そこで，デフォルトの0~9の10種類の数字によるデータセットではなく，0,6の2種類の数字だけのデータセットを生成したいと考えていますが，どうすればいいのかわからず困っています．

該当のソースコード

0~9の10種の数字のデータセットは，以下のコードを実行することで入手することができました．前処理として，データをテンソル型に変換しています．

train_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=True,
    transform=transforms.ToTensor(),
    download=True,
)

train_loader = torch.utils.data.DataLoader(
    dataset=train_dataset,
    batch_size=BATCH_SIZE, #256
    shuffle=True,
    num_workers=0
)

しかし，0,6の2種だけのデータセットはどのように用意すればいいのかわからず見当がつかない状況です．もしご存知の方がいましたら，どうか教えていただけると幸いでございます．どうぞよろしくお願いいたします．

Answer 1

ソースみてみたら、MNIST クラスの属性 self.data, self.targets にデータとラベルをそれぞれ格納しているみたいだったので、それをラベルの種類で制限して上書きすればいいのではないでしょうか。

torchvision.datasets.mnist — PyTorch 1.6.0 documentation

from torch.utils import data
from torchvision import datasets as datasets
from torchvision import transforms

dataset = datasets.MNIST(
    root="/data", train=True, transform=transforms.ToTensor(), download=True,
)

# 中のデータを上書きする
mask = (dataset.targets == 0) | (dataset.targets == 6)
dataset.data = dataset.data[mask]
dataset.targets = dataset.targets[mask]

dataloader = data.DataLoader(
    dataset=dataset, batch_size=32, shuffle=True, num_workers=0
)

# ラベルが制限されていることを確認
for data, label in dataloader:
    print(data.shape)
    print(label)
    break

torch.Size([32, 1, 28, 28])
tensor([6, 0, 0, 6, 0, 6, 6, 6, 6, 0, 0, 0, 6, 0, 0, 6, 0, 0, 0, 6, 6, 6, 0, 0,
        6, 6, 0, 0, 6, 6, 0, 0])

Comments

[Hiro_White03]

tiitoiさん
先ほど無事に，0と6だけのデータセットを入手することができました．ご回答，誠にありがとうございました！