はじめに
この記事ではpytorch-lightningの使い方を、いい感じのプログラムで紹介します
まず結論：ソースコードの書き方
pytorch_lightning_sample.py
import os
import pickle
import numpy as np
from PIL import Image # 画像を取り扱うために使用
import matplotlib.pyplot as plt # 画像のサンプル表示のために使用
import torch # pytorch本体
import torch.nn as nn # ニューラルネットを構成する際の基本的なモジュールが入っている
from torchvision import transforms as transforms # 画像前処理のために使用
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning import loggers as pl_loggers
# ==== データローダの作成など ============================================================================================
# データセットに対して、idxで指定された際に読み込み方法を指定するためのラッパークラス
class Dataset(torch.utils.data.Dataset):
def __init__(self, data, labels):
# 画像を変換して整形して保持
self.data = np.array(data) # numpy形式に変換
self.data = self.data.reshape(len(data), 3, 32, 32) # dataを整形
self.data = self.data.transpose(0, 2, 3, 1) # data[ミニバッチのindex][チャンネル][画像縦位置][画像横位置]と指定できるように順序交換
# ラベルを保持
self.labels = labels
# 画像を前処理するための関数たちを登録
self.transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize(mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5)),
]
)
# 指定されたindexのデータを辞書形式で返却するように設定する
def __getitem__(self, index):
img, label = self.data[index], self.labels[index] # 指定のデータを取得
img = Image.fromarray(img) # 画像に変換
img = self.transform(img) # transformをかける（tensor型に変換してから、正規化）
return {'inputs':img, 'targets':label} # 辞書形式で返却（辞書のkeyはニューラルネットのforwardの引数と同じ名前にする）
# データセットの個数を返すように設定する
def __len__(self):
return len(self.data)
# サンプル表示用の関数。一般には作らなくてOK。
def plot(self, index):
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
img, label = self.data[index], self.labels[index]
plt.imshow(img) # 指定されたindexの画像を描画
plt.title(f'label={classes[label]}') # titleを"label=クラス名"という形式で設定
plt.show() # 表示
class DataModule(pl.LightningDataModule):
def __init__(self, dataset_path, split_ratio=(0.7, 0.1, 0.2), batch_size=32, thread_num=4):
super().__init__()
self.dataset_path = dataset_path
self.split_ratio = split_ratio
self.batch_size = batch_size
self.thread_num = thread_num
def setup(self, stage=None):
# 成形されたデータセットの読み込み
print(f'loading {self.dataset_path}')
with open(self.dataset_path, 'rb') as f:
dataset_raw = pickle.load(f, encoding='bytes')
print('loading completed')
# データセット形式に変換
ds = Dataset(dataset_raw[b'data'], dataset_raw[b'labels'])
ds.plot(index=np.random.randint(0, len(ds))) # サンプルとして、ランダムに1件選んで描画
# データセットを指定された比率に合わせて分割
total_size = len(ds)
train_size = int(total_size * self.split_ratio[0]) # 学習で使用するデータ個数
valid_size = int(total_size * self.split_ratio[1]) # 検証で使用するデータ個数
test_size = int(total_size * self.split_ratio[2]) # テストで使用するデータ個数
self.train_dataset = torch.utils.data.dataset.Subset(ds, range(0, train_size)) # 指定された部分のデータのみを取り出す
self.valid_dataset = torch.utils.data.dataset.Subset(ds, range(train_size, train_size+valid_size)) # 指定された部分のデータのみを取り出す
self.test_dataset = torch.utils.data.dataset.Subset(ds, range(train_size+valid_size, total_size)) # 指定された部分のデータのみを取り出す
print(f'dataset size: total {total_size}, train {train_size}, validation {valid_size}, test {test_size}')
def train_dataloader(self):
if len(self.train_dataset) != 0:
return torch.utils.data.DataLoader(
self.train_dataset, # データセット
batch_size=self.batch_size, # イテレート時のバッチサイズ
shuffle=True, # イテレート前にデータをシャッフルするか
num_workers=self.thread_num, # イテレート時に使用するスレッド数
pin_memory=True, # メモリを固定して高速化をするか
drop_last=True, # 最後の端数部を落とすか
)
else:
raise Exception('length of dataset is zero.')
def val_dataloader(self):
if len(self.valid_dataset) != 0:
return torch.utils.data.DataLoader(
self.valid_dataset, # データセット
batch_size=self.batch_size, # イテレート時のバッチサイズ
shuffle=False, # イテレート前にデータをシャッフルするか
num_workers=self.thread_num, # イテレート時に使用するスレッド数
pin_memory=True, # メモリを固定して高速化をするか
drop_last=False, # 最後の端数部を落とすか
)
else:
raise Exception('length of dataset is zero.')
def test_dataloader(self):
if len(self.test_dataset) != 0:
return torch.utils.data.DataLoader(
self.test_dataset, # データセット
batch_size=self.batch_size, # イテレート時のバッチサイズ
shuffle=False, # イテレート前にデータをシャッフルするか
num_workers=self.thread_num, # イテレート時に使用するスレッド数
pin_memory=True, # メモリを固定して高速化をするか
drop_last=False, # 最後の端数部を落とすか
)
else:
raise Exception('length of dataset is zero.')
# ====================================================================================================================
class LeNet(nn.Module):
def __init__(self):
# ここにはニューラルネットの構成で使用するモジュールを一通り書き出す。
# Conv2dは畳み込み層、AvgPool2dは平均プーリング層、Flattenは数値を一列に並べてベクトル化する層、Linearは線形層（全結合層）を表している。使い方はググろう。
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(3, 6, kernel_size=5, padding=0, stride=1) # 32*32 3チャンネル入力 → 28*28 6チャンネル出力
self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2) # 28*28 6チャンネル入力 → 14*14 6チャンネル出力
self.conv2 = nn.Conv2d(6, 16, kernel_size=5, padding=0, stride=1) # 14*14 6チャンネル入力 → 10*10 16チャンネル出力
self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2) # 10*10 16チャンネル入力 → 5*5 16チャンネル出力
self.flatten = nn.Flatten() # 5*5 25チャンネル入力 → 400(=5*5*25)出力
self.fc1 = nn.Linear(400, 120) # 400入力, 120出力
self.fc2 = nn.Linear(120, 84) # 120入力, 84出力
self.fc3 = nn.Linear(84, 10) # 84入力, 10出力
self.softmax = nn.Softmax(dim=1) # ソフトマックス関数
self.loss = nn.CrossEntropyLoss() # クロスエントロピー損失
def forward(self, inputs, targets):
# ここには順伝搬のやり方を書く。
h = inputs
# 作用
h = torch.sigmoid(self.conv1(h))
h = self.pool1(h)
h = torch.sigmoid(self.conv2(h))
h = self.pool2(h)
h = self.flatten(h)
h = torch.sigmoid(self.fc1(h))
h = torch.sigmoid(self.fc2(h))
h = self.softmax(self.fc3(h))
# 損失を計算
loss = self.loss(h, targets)
# データ形式をdictに包んで出力
return {
"loss":loss, # 損失の値を記録
"hidden_states":h, # 最終層の値を記録
}
# ====================================================================================================================
class LitModule(pl.LightningModule):
# ネットワークモジュールなどの定義
def __init__(self, learning_rate):
super().__init__()
self.model = LeNet()
self.learning_rate = learning_rate
self.save_hyperparameters()
# 順伝搬の処理
def forward(self, **x):
return self.model(**x)
# オプティマイザの定義
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1, end_factor=0.1, total_iters=self.trainer.max_epochs)
return [optimizer], [scheduler]
# ==================================================================
# 学習のバッチ実行処理
def training_step(self, batch, batch_index):
outputs = self.model(**batch)
loss = outputs['loss']
return {'loss': loss, 'correct': (torch.argmax(outputs["hidden_states"], dim=-1)==batch['targets']).to(torch.float32)}
# 学習の全バッチ終了時の処理
def training_epoch_end(self, outputs):
train_loss = torch.hstack([dict_['loss'] for dict_ in outputs]).mean()
train_accuracy = torch.hstack([dict_['correct'] for dict_ in outputs]).mean()
self.log_dict({"train_loss": train_loss, "train_accuracy": train_accuracy})
self.print({"train_loss": train_loss, "train_accuracy": train_accuracy})
# 検証のバッチ実行処理
def validation_step(self, batch, batch_index):
with torch.no_grad():
outputs = self.model(**batch)
loss = outputs['loss']
return {'loss': loss, 'correct': (torch.argmax(outputs["hidden_states"], dim=-1)==batch['targets']).to(torch.float32)}
# 検証の全バッチ終了時の処理
def validation_epoch_end(self, outputs):
val_loss = torch.hstack([dict_['loss'] for dict_ in outputs]).mean()
val_accuracy = torch.hstack([dict_['correct'] for dict_ in outputs]).mean()
self.log_dict({"val_loss": val_loss, "val_accuracy": val_accuracy})
self.print({"val_loss": val_loss, "val_accuracy": val_accuracy})
# テストのバッチ実行処理
def test_step(self, batch, batch_index):
with torch.no_grad():
outputs = self.model(**batch)
loss = outputs['loss']
return {'loss': loss, 'correct': (torch.argmax(outputs["hidden_states"], dim=-1)==batch['targets']).to(torch.float32)}
# テストの全バッチ終了時の処理
def test_epoch_end(self, outputs):
test_loss = torch.hstack([dict_['loss'] for dict_ in outputs]).mean()
test_accuracy = torch.hstack([dict_['correct'] for dict_ in outputs]).mean()
self.print({"test_loss": test_loss, "test_accuracy": test_accuracy})
# ==================================================================
# ====================================================================================================================
def main():
save_dir = './result'
data_module = DataModule(
dataset_path='./cifar-10-batches-py/data_batch_1',
split_ratio=(0.7, 0.1, 0.2),
batch_size=32,
)
model = LitModule(learning_rate=0.0001)
callbacks = [
ModelCheckpoint(
dirpath=save_dir,
filename='epoch{epoch:02d}-val_loss{val_loss:.2f}', # チェックポイントのファイル名の形式
monitor='val_loss', # 基準とする量
mode="min", # 最小となるところを探す
), # modelのチェックポイントを作成
# EarlyStopping(
# monitor="val_loss", # 基準とする量
# mode="min", # 最小となるところを探す
# ), # early-stoppingを利用
]
trainer = pl.Trainer(
max_epochs=10,
logger=[pl_loggers.TensorBoardLogger(save_dir=save_dir)],
callbacks=callbacks,
accelerator='gpu',
devices=[0], # 使用するGPUのIDのリスト
# auto_lr_find=True, # learning rateを自動で設定するか
# accumulate_grad_batches=1, # 勾配を累積して一度に更新することでバッチサイズを仮想的にN倍にする際のN
# gradient_clip_val=1, # 勾配クリッピングの値
# fast_dev_run=True, # デバッグ時にonにすると、1回だけtrain,validを実行する
# overfit_batches=1.0, # デバッグ時にonにすると、train = validで学習が進み、過学習できているかを確認できる
# deterministic=True, # 再現性のために乱数シードを固定するか
# resume_from_checkpoint='bbb/aaa.ckpt', # チェックポイントから再開する場合に利用
# precision=16, # 小数を何ビットで表現するか
# amp_backend="apex", # 少数の混合方式を使用するかどうか。nvidiaのapexがインストールされている必要あり。
# benchmark=True, # cudnn.benchmarkを使用して高速化するか（determisticがTrueの場合はFalseに上書きされる）
)
# trainer.tune(model, datamodule=data_module) # 「auto_lr_find=True」を指定した場合に実行する
trainer.fit(model, datamodule=data_module)
trainer.test(model, datamodule=data_module)
if __name__ == '__main__':
main()
終わりに
- とりあえずメモ書き程度で残しておきます
- 今後もう少し説明を書くかも