[Pytorch]DatasetとDataloaderのイメージ図（自作Dataset版）

目的

『DatasetとDataloaderのイメージ図』でMNISTに用いるDatasetとDataloaderのイメージ図を作成した。理解を深めるために単純なMLPを用いた単純な回帰問題を自作のDatasetとDataloaderを用いて解いたので共有する。

データ処理のイメージ図と抜粋版スクリプト

■ データの生成　→　学習用と検証用データに分割　→　Dataset作成　→　Dataloader作成　→　Dataloaderからミニバッチの取り出し部分　の部分のスクリプトを抜粋。
■ DatasetとDataloaderの理解に邪魔な部分は削除していますので、実際にスクリプトを動かしてみたい方は後述の『実際に動作するスクリプト』を使用してください。

■ データ処理のイメージ図
学習用と検証用データに同じ処理をしているのでイメージ図は学習用データのみとなっています。

抜粋版スクリプト.py

#---------------データの生成-------------------------
data_size = 1000 #後で学習用500, 検証用500に分割する
x = torch.randn(data_size, 1)  # 正規分布に従う乱数
y = 2 * x + 3 + torch.randn(data_size, 1) * 0.3  # 動作確認時はノイズの係数を0にすると見やすい

#---------------データを学習用と検証用に分割---------------
train_size = int(0.5 * len(x))
x_train, x_vali = torch.split(x, [train_size, len(x) - train_size])
y_train_act, y_vali_act = torch.split(y, [train_size, len(y) - train_size])

#---------自作データセットクラスの定義---------
class CustomDataset(Dataset):
    def __init__(self, x_data, y_data):
        self.x_data = x_data
        self.y_data = y_data
    
    def __len__(self):
        return len(self.x_data)
    
    def __getitem__(self, idx):
        return self.x_data[idx], self.y_data[idx]
        
#---------データセットのインスタンス化---------
dataset_train = CustomDataset(x_train, y_train_act)
dataset_vali = CustomDataset(x_vali, y_vali_act)

#---------DataLoaderのインスタンス作成----------
batch_size = 100
loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=False)
loader_vali = DataLoader(dataset_vali, batch_size=batch_size, shuffle=False)

#----------ミニバッチごとにデータを取り出してサイズの確認する-------
 for i, (x_batch, y_batch) in enumerate(loader_train):
     print('Batch :', i, ',   x_batch.size:', x_batch.size(), ',   y_batch.size:', y_batch.size())

実際に動作するスクリプト

d13d_Pytorch_MLP_Simple_Regression_with_Dataloader.py

#---------------モジュールのインポート-------------------------
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
torch.manual_seed(0)

#---------------データの生成-------------------------
data_size = 1000
x = torch.randn(data_size, 1)  # 正規分布に従う乱数
y = 2 * x + 3 + torch.randn(data_size, 1) * 0.3  # 動作確認時はノイズの係数を0にすると見やすい

#print(x.size())
#print(y.size())

#---------------データを学習用と検証用に分割---------------
train_size = int(0.5 * len(x))
x_train, x_vali = torch.split(x, [train_size, len(x) - train_size])
y_train_act, y_vali_act = torch.split(y, [train_size, len(y) - train_size])

#print(x_train.size())
#print(y_train_act.size())

#---------自作データセットクラスの定義---------
class CustomDataset(Dataset):
    def __init__(self, x_data, y_data):
        self.x_data = x_data
        self.y_data = y_data
    
    def __len__(self):
        return len(self.x_data)
    
    def __getitem__(self, idx):
        return self.x_data[idx], self.y_data[idx]

#---------データセットのインスタンス化---------
dataset_train = CustomDataset(x_train, y_train_act)
dataset_vali = CustomDataset(x_vali, y_vali_act)

#---------DataLoaderのインスタンス作成----------
batch_size = 100
loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=False)
loader_vali = DataLoader(dataset_vali, batch_size=batch_size, shuffle=False)

#----------ミニバッチごとにデータを取り出してサイズの確認する-------
for i, (x_batch, y_batch) in enumerate(loader_train):
    print('Batch :', i, ',   x_batch.size:', x_batch.size(), ',   y_batch.size:', y_batch.size())

#---------------MLPモデルの定義---------------
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(1, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

#---------------モデルのインスタンス化------
model = MLP()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

#---------------モデルのミニバッチ学習---------------
num_epochs = 500
for epoch in range(num_epochs):
    loss_epoch = 0 #1epochの損失
    for x_batch, y_batch in loader_train:
        model.train()
        optimizer.zero_grad()
        outputs = model(x_batch)
        loss = criterion(outputs, y_batch)
        loss_epoch = loss_epoch + loss
        loss.backward()
        optimizer.step()
    
    if (epoch+1) % 100 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss_epoch:.4f}')

#---------------学習用データに対する予測(ミニバッチ毎に実施)-----------------
model.eval()
y_train_pred = []
with torch.no_grad():
    for x_batch, _ in loader_train:
        outputs = model(x_batch)
        y_train_pred.append(outputs)
        
#print(type(y_train_pred))
y_train_pred = torch.cat(y_train_pred, dim=0) #listをtorch.tensorに変換する
#print(type(y_train_pred))


#---------------検証用データに対する予測(ミニバッチ毎に実施)-----------------
model.eval()
y_vali_pred = []
with torch.no_grad():
    for x_batch, _ in loader_vali:
        outputs = model(x_batch)
        y_vali_pred.append(outputs)

#print(type(y_vali_pred))
y_vali_pred = torch.cat(y_vali_pred, dim=0) #listをtorch.tensorに変換する
#print(type(y_vali_pred))

#---------------予測結果のプロット------------------------------
plt.figure(figsize=(5, 5))
plt.scatter( y_train_act.numpy(), y_train_pred.numpy(), label='Train Data')
plt.scatter( y_vali_act.numpy(), y_vali_pred.numpy(), label='Vali Data')
plt.plot([-10, 10], [-10, 10], color='gray', label='1:1 line')
plt.title('Prediction result')
plt.xlabel('Y act')
plt.ylabel('Y pred')
plt.grid(True, linestyle='dotted')
plt.legend()

plt.show()

■ 学習用データと予測データに対する検証結果