人気急上昇PyTorchで学ぶディープラーニング入門【データセット・データローダー編】

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import math
%matplotlib inline

class Model(nn.Module):
    def __init__(self, in_features=4, h1=8, h2=9, out_features=3):
        super().__init__()
        self.fc1 = nn.Linear(in_features, h1)
        self.fc2 = nn.Linear(h1, h2)
        self.out = nn.Linear(h2, out_features)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.out(x)
        return x

torch.manual_seed(1)
model = Model()
# 損失関数
criterion = nn.CrossEntropyLoss()
# 最適化関数
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# データセット定義
class IrisDataset(Dataset):
    
    def __init__(self, is_train=True):
        # データのロード
        iris = load_iris()
        df = pd.DataFrame(iris.data, columns=iris.feature_names)
        # 花の種類（target）のカラムを作成
        df['target'] = iris.target
        
        # 入力データの準備
        X = df.drop('target', axis=1).values
        y = df['target'].values
        # 学習データと検証データに分割
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=123)
        
        if is_train:
            self.X_data = torch.tensor(X_train, dtype=torch.float32)
            self.y_data = torch.tensor(y_train)
        else:
            self.X_data = torch.tensor(X_test, dtype=torch.float32)
            self.y_data = torch.tensor(y_test)
        
        self.data_length = len(self.X_data)
        
    def __getitem__(self, index):
        return self.X_data[index], self.y_data[index]
        
    def __len__(self):
        return self.data_length

BATCH_SIZE = 20
dataset = IrisDataset()
train_data_loader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, shuffle=True)

epochs = 100
loss_list = []
total_samples = len(dataset)
# データ件数
print(total_samples)
n_iterations = math.ceil(total_samples / BATCH_SIZE)
# イテレーションの数
print(n_iterations)

for epoch in range(epochs):
    print(f'Epoch: {epoch+1}/{epochs}')
    # バッチサイズごとにループ
    for i, data in enumerate(train_data_loader):
        inputs, labels = data
        print(f'Iteration: {i+1}/{n_iterations}, Inputs: {inputs.shape}, Labels: {labels.shape}')
        y_pred = model.forward(inputs.data)
        loss = criterion(y_pred, labels.data)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        loss_list.append(loss)


# Epoch: 1/100
# Iteration: 1/6, Inputs: torch.Size([20, 4]), Labels: torch.Size([20])
# Iteration: 2/6, Inputs: torch.Size([20, 4]), Labels: torch.Size([20])
# Iteration: 3/6, Inputs: torch.Size([20, 4]), Labels: torch.Size([20])
# Iteration: 4/6, Inputs: torch.Size([20, 4]), Labels: torch.Size([20])
# Iteration: 5/6, Inputs: torch.Size([20, 4]), Labels: torch.Size([20])
# Iteration: 6/6, Inputs: torch.Size([20, 4]), Labels: torch.Size([20])

plt.plot(loss_list)