pytorchで自作データセットを用いる方法

ここでは、pytorchで用意されているデータセット以外を用いる際のデータローダの作成方法を解説します。
今回はsklearnで用意されているデータセットを用います。

実装

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torch import optim
from torch.utils.data import DataLoader, Dataset, TensorDataset

# MNISTデータセットの読み込み
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split

mnist_X, mnist_y = fetch_openml('mnist_784', version=1, data_home="./mnist", return_X_y=True)

X_all = mnist_X.astype(np.float32) / 255 # データを[0,1]に変換
y_all = mnist_y.astype(np.int32)

# 分割
X_train, X_test, y_train, y_test = train_test_split(X_all, y_all, train_size=0.8, random_state=1)

print(X_train.shape) # (56000, 784)

X_train = X_train.reshape(len(X_train), 1, 28, 28).astype('float32')
X_test = X_test.reshape(len(X_test), 1, 28, 28).astype('float32')

y_train = y_train.astype('long')
y_test = y_test.astype('long')

train_set = TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train))
test_set  = TensorDataset(torch.from_numpy(X_test), torch.from_numpy(y_test))

bs = 1000

train_loader = DataLoader(train_set, batch_size=bs, shuffle=True)
test_loader = DataLoader(test_set, batch_size=bs, shuffle=False)

解説

train_set = TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train))
test_set  = TensorDataset(torch.from_numpy(X_test), torch.from_numpy(y_test))

ここでテンソルデータセットを作成しているが、入力画像データ(X)は (データ数、チャネル数、高さ、幅)に設定する必要がある。
そのため、

X_train = X_train.reshape(len(X_train), 1, 28, 28).astype('float32')
X_test = X_test.reshape(len(X_test), 1, 28, 28).astype('float32')

で変換している。
また、正解ラベルは

y_train = y_train.astype('long')
y_test = y_test.astype('long')

のようにlong型に設定する必要がある。