pytorchのお勉強 ～ゼロから作るpytorch-lightning での線形回帰～

0. 概要

• pytorch及びpytorch-lightningを使った機械学習の基本的な書き方の備忘録

• 次のステップで線形回帰のプログラムを書き換えた記録

  1. pytorchなしで線形回帰
  2. pytorchで線形回帰　その1 (loss, optimizerの利用)
  3. pytorchで線形回帰　その2 (modelのclass定義とdataset, dataloaderの利用)
  4. pytorch-lightningで線形回帰

• チュートリアル見ながらpytorchでコードを書いたけど、理解不十分でもやもやしていたのが少し解消できた。同様の悩みを持っている方の助けになれば幸いです。

• torchとlightningは次のversionを使って実行しました。

torch 1.5.1
pytorch-lightning 1.3.6

1. コード最終形態

こちらのコードを今回のゴールにします。

# %%========================================================
# ライブラリの読込み
# ==========================================================
import argparse
import torch
import matplotlib.pyplot as plt
import pytorch_lightning as pl

# %%========================================================
# 学習用データの準備とグラフ化表示
# ==========================================================
# 学習用データの準備
# 適当に作ったx、yの組合せを予測できるモデルを学習する

X = [0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]
Y = [0.1, 0.15, 0.16,0.12, 0.15,0.16, 0.2, 0.24, 0.16,0.17, 0.28,0.25, 0.3, 0.24, 0.31,0.35, 0.3, 0.4,  0.41,0.32,0.42]

# 学習データをプロットして表示
# fig = plt.figure()
# plt.scatter(X, Y)
# plt.show()

# %%========================================================
# model class
# ==========================================================
# torch.nn.Moduleを継承して書くこと
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Parameter(torch.zeros(1))
        self.b = torch.nn.Parameter(torch.zeros(1))
    
    def forward(self, x):
        return self.a * x + self.b


# %%========================================================
# dataset class
# ==========================================================
# torch.utils.data.Datasetを継承して作る
class Dataset(torch.utils.data.Dataset):
    def __init__(self, X, Y):
        super().__init__()
        self.x = torch.tensor(X)
        self.y = torch.tensor(Y)

    def __len__(self):
        return self.x.shape[0]

    def __getitem__(self, index):
        return self.x[index], self.y[index]


# %%========================================================
# lightning data module
# ==========================================================
class LitModule(pl.LightningDataModule):
    def __init__(self, X, Y, batch_size):
        super().__init__()

        # 変更する可能性の変数は外から指定できるように
        self.batch_size = batch_size
        # datasetを作成
        self.setup_dataset(X, Y)

    def setup_dataset(self, x, y):
        torch.manual_seed(10)
        dataset = Dataset(x, y)
        self.dataset_train, self.dataset_valid, self.dataset_test = \
                    torch.utils.data.random_split(dataset, [13, 4, 4])

    # この通りの名前でメソッドを定義するとlightningが自動で使い分けてくれる
    def train_dataloader(self):
        return torch.utils.data.DataLoader(self.dataset_train, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return torch.utils.data.DataLoader(self.dataset_valid, batch_size=self.batch_size, shuffle=False)

    def test_dataloader(self):
        return torch.utils.data.DataLoader(self.dataset_test, batch_size=self.batch_size, shuffle=False)


# %%========================================================
# lightning model
# ==========================================================
class LitModel(pl.LightningModule):
    def __init__(self, args):
        super().__init__()
        
        self.save_hyperparameters(args)  # hparams.yamlに自動で書き出してくれる
        self.model = Model()
        self.loss_fn = torch.nn.MSELoss()
        self.optimizer = torch.optim.SGD(self.parameters(), lr=args.a1_lr)

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        loss = self.shared_step(batch)
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
        return loss

    def validation_step(self, batch, batch_idx):
        loss = self.shared_step(batch)
        self.log('valid_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)

    def shared_step(self, batch):
        xi, yi = batch
        output  =self(xi)
        loss = self.loss_fn(output, yi)
        return loss

    # def validation_epoch_end(self, validation_step_outputs):
    #     for out in validation_step_outputs:

    def configure_optimizers(self):
        return self.optimizer


# %%========================================================
# 学習
# ==========================================================
# ハイパーパラメータの設定
parser = argparse.ArgumentParser(description="linear Regression")
parser.add_argument("--a1_lr",            type=float,   default=0.01)
parser.add_argument("--a2_batch_size",    type=int,     default=2)
parser.add_argument("--a3_max_epochs",    type=int,     default=50)
parser.add_argument("--a4_gpus",          type=int,     default=0)
parser.add_argument("--d1_train_data_x",  type=float,  nargs="+",  default=X)
parser.add_argument("--d2_train_data_y",  type=float,  nargs="+",  default=Y)

args = parser.parse_args()

pl_model = LitModel(args)
pl_module = LitModule(args.d1_train_data_x, args.d2_train_data_y, batch_size=args.a2_batch_size)

early_stopping = pl.callbacks.EarlyStopping("valid_loss")
model_chkpoint = pl.callbacks.ModelCheckpoint(save_last=True, save_top_k=2, monitor="valid_loss")
trainer = pl.Trainer(gpus=args.a4_gpus, max_epochs=args.a3_max_epochs, callbacks=[early_stopping, model_chkpoint])
trainer.fit(pl_model, pl_module)

print("training finishied")
print(pl_model.logger.log_dir)

2. pytorchなしで線形回帰

###2-1. 題材とする学習データ

# %%========================================================
# ライブラリの読込み
# ==========================================================
# データ確認用にmatplotlibを読込み
import matplotlib.pyplot as plt

# %%========================================================
# 学習用データの準備とグラフ化表示
# ==========================================================
# 学習用データの準備
# 適当に作ったx、yの組合せを予測できるモデルを学習する
X = [0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]
Y = [0.1, 0.15, 0.16,0.12, 0.15,0.16, 0.2, 0.24, 0.16,0.17, 0.28,0.25, 0.3, 0.24, 0.31,0.35, 0.3, 0.4,  0.41,0.32,0.42]

# 学習データをプロットして表示
fig = plt.figure()
plt.scatter(X, Y)
plt.show()

適当に作った下図のプロットを線形回帰することを題材にする

###2-2.学習部分
y = ax + bで近似することにして、初期値をa=0, b=0に設定
２乗誤差が小さくなるように勾配降下法でa,bを少しずつ修正する

# %%========================================================
# 線形回帰　pytorchなし
# ==========================================================
# ハイパーパラメータ（自分で決めるパラメータ）の設定
epochs = 500
lr = 0.01

# a,bの初期化。今回は簡単のため0とする。
a = 0
b = 0

# 学習
# 精度が出るまでepochs回繰り返す
for epoch in range(1, epochs+1):
    loss_iter = 0
    # x, y を一つずつ取り出して、計算して、2乗誤差が小さくなるようa、bを少しずつ修正していく
    for xi, yi in zip(X, Y):
        output = a * xi + b        # y = ax + bを計算
        loss1 = (output - yi)      # 正解との差を算出。逆伝播に使うので2乗する前を保管
        loss2 = loss1 ** 2         # 2乗誤差
        # 簡単な式なので、逆伝播も簡単だが、複雑モデルになるとすごく大変
        a -= loss1 * xi * lr       # aを更新　逆伝播　
        b -= loss1 * lr            # bを更新　逆伝播
        loss_iter += loss2         # epoch全体での誤差を算出
    if epoch % (epochs // 20) == 0:
        print("epoch={:03d} | a={:.3f}, b={:.3f} | loss_iter={:.3f}".format(epoch, a, b, loss_iter))

上記コードを実行すると、徐々に誤差が小さくなる様子が確認できる

epoch=015 | a=0.143, b=0.181 | loss_iter=0.073
epoch=030 | a=0.176, b=0.167 | loss_iter=0.056
epoch=045 | a=0.200, b=0.153 | loss_iter=0.046
epoch=060 | a=0.220, b=0.142 | loss_iter=0.039
...(略)
epoch=300 | a=0.294, b=0.101 | loss_iter=0.027

###2-3. 学習結果の確認

# 学習結果の確認
fig = plt.figure()
y_predict = [a*xi+b for xi in X]
plt.scatter(X, Y, c="blue", label="training data")
plt.plot(X, y_predict, c="red", label="model prediction")
plt.legend()
plt.show()

青点が学習データ、赤線が学習後のモデル(y=0.294x+0.101)
回帰直線が引けた

#3. pytorchで線形回帰　その1 (loss, optimizer)
###3-1.学習

# %%========================================================
# 線形回帰　with pytorch(loss, optimizer) 
# ==========================================================
# 学習データをtorch tensorに変換
x_t = torch.tensor(X)
y_t = torch.tensor(Y)

# ハイパーパラメータの設定
epochs = 300
lr = 0.01

# a,bの初期化。前回同様0。1変数だけど、torch tensorにする
# requirse_grad = Trueにすることで逆伝播の計算を自動でやってくれるようになる
a = torch.zeros(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)

# 誤差の計算と変数の更新はpytorchにある機能を使う
# lossの計算をしてくれる関数
loss_fn = torch.nn.MSELoss()
# optimizerを設定すると面倒な誤差逆伝播をpytorchがやってくれる     
optimizer = torch.optim.SGD((a,b), lr=lr) # 第一引数に更新するパラメータを渡す。イテラブルな形式で渡す

# 学習
for epoch in range(1, epochs+1):
    loss_iter = 0
    for xi, yi in zip(x_t, y_t):
        output = a * xi + b
        loss = loss_fn(output, yi)  # ①目的変数の計算
        optimizer.zero_grad()       # ②各テンソルが持っている勾配の値を初期化
        loss.backward()             # ③誤差逆伝播。自動で計算してくれる
        optimizer.step()            # ④各変数の値を更新する
        loss_iter += loss
    if epoch % (epochs // 20) == 0:
        print("epoch={:03d} | a={:.3f}, b={:.3f} | loss_iter={:.3f}".format(epoch, a.item(), b.item(), loss_iter))

①誤差を求めて
②各テンソルが裏で持ってる勾配情報を初期化してから
③逆伝播（各変数の微分を算出）
④変数をアップデート

• ②を忘れると正しく微分係数が求まらない（足し合わされる？）ので注意
• 逆伝播を自動でやってくれるので、微分の逆伝播計算に不安がある私のような人も安心
  (線形回帰のコードだと簡単になったように見えないが、deepになると自分で書くのは難しい）

###3-2.学習結果の確認

fig = plt.figure()
y_predict = [a.item()*xi+b.item() for xi in X]
plt.scatter(X, Y, c="blue", label="training data")
plt.plot(X, y_predict, c="red", label="model prediction")
plt.legend()
plt.show()

pytorchを使わない場合と同様に回帰直線が引けた
x,yがtorch.tensorになっているので、.item() もしくは、.detach().numpy()
で変数かnumpy形式にすることが必要(matplotのversionによる？）

#4. pytorchで線形回帰　その2 (modelのclass定義とdataset, dataloaderの利用)
###4-1. modelをclassで定義する
線形回帰でなくdeep learningをする場合はモデルが複雑になるので、classにして独立させた方が管理しやすい。

# torch.nn.Moduleを継承して書く
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Parameter(torch.zeros(1))
        self.b = torch.nn.Parameter(torch.zeros(1))
    
    def forward(self, x):
        return self.a * x + self.b

class定義したモデルで学習する

model = Model()

# nn.Muduleを継承して書けば、.parameters()でモデルの重みを取得できるので、
# model.parameters()をoptimizerに渡す
optimizer = torch.optim.SGD(model.parameters(), lr)

# modelを使って計算するときはforward()でなく、call()を使う
# pytorchがいい感じに処理して、forwardで定義した内容を実行してくれるらしい
output = model(xi)

###4-2. datasetとdataloaderを活用する

ここまでは学習データから一つずつ値を取り出して(xi, yi)予測値の計算⇒誤差計算⇒逆伝播していた。

for xi, yi in zip(x_t, y_t):

このように一つずつ取り出して学習することをオンライン学習と呼ぶ。
計算効率向上などの目的で、ある大きさの塊（ミニバッチ）に分けて学習させる”ミニバッチ学習”が現在の主流。
このミニバッチを自動で作ってくれるのが、dataloaderで以下のコードで簡単に実装可能

dataloader = torch.utils.data.DataLoader(データ, batch_size=2, shuffle=True)

もちろん学習用データを引数で渡す必要があるが、dataloaderが内部で、

• 全データ数を調べる "len(データ）"
• 指定したidのデータを抜き出す "データ[id]"
  を実行するので、その機能を持っている必要がある。(例えばlistはその条件を満たす）

通常はDataset classを作り、dataloaderに渡す

# torch.utils.data.Datasetを継承して作る
class Dataset(torch.utils.data.Dataset):
    def __init__(self, X, Y):
        super().__init__()
        self.x = torch.tensor(X)
        self.y = torch.tensor(Y)

    # __len__(self) で len()で返す値を指定できる
    def __len__(self):
        return self.x.shape[0]

    # __getitem__(self, index)で index番目の結果を返す処理を定義できる
    def __getitem__(self, index):
        return self.x[index], self.y[index]

datasetの動作確認

>>> dataset = Dataset(X, Y)
>>> print("len()でデータの数が取得できること。len(dataset)=", len(dataset))
len()でデータの数が取得できること。len(dataset)= 21
>>> print("[index]でindex番目のデータを取り出せること。dataset[5]=",dataset[5])
[index]でindex番目のデータを取り出せること。dataset[5]= (tensor(0.2500), tensor(0.1600))

datasetを自動で分割してくれる機能もあるので、trainとvalidに分割することとする

# dataset を trainとvalidationに分ける
# 毎回同じように分割させるためにシードを固定
torch.manual_seed(10)
dataset_train, dataset_valid = torch.utils.data.random_split(dataset, [16, 5]) # [16,5]はtrain,validの個数

# 作成したDatasetをDataLoaderに渡せばdataloaderが完成する
dataloader_train = torch.utils.data.DataLoader(dataset_train, batch_size=2, shuffle=True)
dataloader_valid = torch.utils.data.DataLoader(dataset_valid, batch_size=2, shuffle=False)

# dataset, dataloaderを使うことで、
# ランダムに分割したtrain, validデータセットをバッチサイズ毎に取り出すことができる
# trainはランダムな組み合わせ&順番で、validは固定した順番で取り出す
print("train")
for x, y in dataloader_train:
    print(x, y)
print("valid")
for x, y in dataloader_valid:
    print(x, y)

上記を実行すると、下記のようにバッチサイズ２のテンソルが返ってくる。バッチの組み合わせもランダムになっている。

train
tensor([0.1000, 0.4000]) tensor([0.1600, 0.1600])
tensor([0.7500, 0.3500]) tensor([0.3500, 0.2400])
tensor([0.3000, 0.2000]) tensor([0.2000, 0.1500])
tensor([0.0500, 0.5000]) tensor([0.1500, 0.2800])
tensor([0.1500, 0.2500]) tensor([0.1200, 0.1600])
tensor([0.6500, 0.7000]) tensor([0.2400, 0.3100])
tensor([1.0000, 0.8000]) tensor([0.4200, 0.3000])
tensor([0.9000, 0.5500]) tensor([0.4100, 0.2500])
valid
tensor([0.6000, 0.9500]) tensor([0.3000, 0.3200])
tensor([0.8500, 0.4500]) tensor([0.4000, 0.1700])
tensor([0.]) tensor([0.1000])

注意）seed固定しているので同じ結果になる。epochが変われば、trainのバッチの組み合わせが変わる

###4-3.modelのclass定義、dataset, dataloader利用版まとめ
（コードは上述と重複しています）

# %%========================================================
# ライブラリの読込み
# ==========================================================
import random
import torch
import matplotlib.pyplot as plt

# %%========================================================
# 学習用データの準備とグラフ化表示
# ==========================================================
# 学習用データの準備
# 適当に作ったx、yの組合せを予測できるモデルを学習する

X = [0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]
Y = [0.1, 0.15, 0.16,0.12, 0.15,0.16, 0.2, 0.24, 0.16,0.17, 0.28,0.25, 0.3, 0.24, 0.31,0.35, 0.3, 0.4,  0.41,0.32,0.42]

# 学習データをプロットして表示
fig = plt.figure()
plt.scatter(X, Y)
plt.show()

# %%========================================================
# model class
# ==========================================================
# torch.nn.Moduleを継承して書くこと
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Parameter(torch.zeros(1))
        self.b = torch.nn.Parameter(torch.zeros(1))
    
    def forward(self, x):
        return self.a * x + self.b

# %%========================================================
# dataset class
# ==========================================================
# torch.utils.data.Datasetを継承して作る
class Dataset(torch.utils.data.Dataset):
    def __init__(self, X, Y):
        super().__init__()
        self.x = torch.tensor(X)
        self.y = torch.tensor(Y)

    def __len__(self):
        return self.x.shape[0]

    def __getitem__(self, index):
        return self.x[index], self.y[index]

# %%========================================================
# 線形回帰 with pytorch (dataset, dataloader, model)
# ==========================================================

# ハイパーパラメータの設定
epochs = 100
lr = 0.01

# 学習データの準備(dataset, dataloader)、モデルの呼び出し
dataset = Dataset(X, Y)
torch.manual_seed(10)
dataset_train, dataset_valid = torch.utils.data.random_split(dataset, [16, 5])
dataloader_train = torch.utils.data.DataLoader(dataset_train, batch_size=2, shuffle=True)
dataloader_valid = torch.utils.data.DataLoader(dataset_valid, batch_size=2, shuffle=False)

model = Model()
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

# 学習
loss_logger_train = []
loss_logger_valid = []
for epoch in range(1, epochs+1):
    # train step
    loss_iter = 0
    for xi, yi in dataloader_train:
        # 推測と誤差算出（validと共通）
        output = model(xi)
        loss = loss_fn(output, yi)
        loss_iter += loss
        # 逆伝播　train時のみ
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    loss_logger_train.append(loss_iter)

    # valid step
    loss_iter = 0
    for xi, yi in dataloader_valid:

        output = model(xi)
        loss = loss_fn(output, yi)
        loss_iter += loss

    loss_logger_valid.append(loss_iter)

    if epoch % (epochs // 20) == 0:
        print("epoch={:03d} | train_loss={:.3f} | valid_loss={:.3f}".format(epoch, loss_logger_train[-1].item(), loss_logger_valid[-1].item()))

# グラフ化
model.eval()
y_predict = [model(xi).item() for xi in X]
fig, axes = plt.subplots(1,2, figsize=(12,5))
ax = axes.ravel()
ax[0].scatter(X, Y, c="blue", label="training data")
ax[0].plot(X, y_predict, c="red", label="model prediction")
loss_x = [i for i in range(epochs)]
ax[1].plot(loss_x, loss_logger_train, label="train_loss")
ax[1].plot(loss_x, loss_logger_valid, label="valid_loss")
plt.legend()
plt.show()

実行結果

epoch=005 | train_loss=0.118 | valid_loss=0.040
epoch=010 | train_loss=0.041 | valid_loss=0.019
…(略)
epoch=095 | train_loss=0.012 | valid_loss=0.007
epoch=100 | train_loss=0.011 | valid_loss=0.007

• これまで同様に回帰直線が引けた（左）
• trainとvalidに分けたので誤差の推移グラフを描いて汎化性能の確認が可能になった（右）

#5. pytorch-lightningで線形回帰
ライブラリの読込み

# ライブラリの読込み
import pytorch_lightning as pl

pytorch-lightningで書く場合の変更点と変わらない点

• model, datasetは上記を同じものを使用する
• dataloaderはlightning Module classに記述する
• 学習部分はlightning Model classに記述する
• 学習条件はtrainerに設定して、trainer.fit()で学習する

###5-1. lightning Module

datasetを渡すか、内部で定義して、dataloaderを定義する

• def train_dataloader(self):
• def val_dataloader(self):
  にそれぞれtrain, validのdataloaderを定義することで、lightningが学習と検証時で使い分けてくれる

# %%========================================================
# lightning data module
# ==========================================================
# pl.LightningDataModuleを継承する
class LitModule(pl.LightningDataModule):
    def __init__(self, X, Y, batch_size):
        super().__init__()

        # 変更する可能性の変数は外から指定できるようにしておく
        self.batch_size = batch_size
        # datasetを作成
        self.setup_dataset(X, Y)

    # datasetの生成メソッド
    def setup_dataset(self, x, y):
        torch.manual_seed(10)
        dataset = Dataset(x, y)
        self.dataset_train, self.dataset_valid, self.dataset_test = \
                    torch.utils.data.random_split(dataset, [13, 4, 4])

    # この通りの名前でメソッドを定義するとlightningが自動で使い分けてくれる
    def train_dataloader(self):
        return torch.utils.data.DataLoader(self.dataset_train, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return torch.utils.data.DataLoader(self.dataset_valid, batch_size=self.batch_size, shuffle=False)

    def test_dataloader(self):
        return torch.utils.data.DataLoader(self.dataset_test, batch_size=self.batch_size, shuffle=False)

###5-2. lightning model

# %%========================================================
# lightning model
# ==========================================================
# pl.LightningModuleを継承する
class LitModel(pl.LightningModule):
    def __init__(self, hp):
        super().__init__()

        # hparams.yamlに自動で書き出してくれる
        # lightning の versionによって書き方が違う模様
        self.save_hyperparameters(hp)  

        self.model = Model()
        self.loss_fn = torch.nn.MSELoss()
        self.optimizer = torch.optim.SGD(self.parameters(), lr=hp["11.lr"])

    # forwardを実行するとモデルの結果が返るように設定する
    def forward(self, x):
        return self.model(x)

    # training stepで実施することを定義する
    # optimizerは自動でやってくれるので、returnでlossを返すところまで
    # 引数batchはdataloaderによって渡されるX, yのミニバッチ
    def training_step(self, batch, batch_idx):
        loss = self.shared_step(batch)
        # self.logでログを記録できる。デフォルトはtensorboard形式
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
        return loss

    # validationもtraining同様だがreturnは不要
    def validation_step(self, batch, batch_idx):
        loss = self.shared_step(batch)
        self.log('valid_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)

    # train, validで共通な部分は独立したメソッドにしておくと、修正する時が楽
    def shared_step(self, batch):
        xi, yi = batch
        # forwardはself()で実行できる。self.forward(xi)とはしない
        output  =self(xi)
        loss = self.loss_fn(output, yi)
        return loss

    # optimizer.step()などはlightningが自動でしてくれる
    # configure_optimizersで使用するoptimizerを返すだけでOK
    def configure_optimizers(self):
        return self.optimizer

以上で、lightning module と lightning modelの設定は終了
pytorch-lightningにしても学習部分コードの量は減らないが、
整然としているので、チームで共有化するときに向いていると感じる。
また、self.logでtensorboard用のファイルが出力されるので、出力処理がすっきり書ける。

###5-3. 学習の実行

# %%========================================================
# 学習
# ==========================================================
# ハイパーパラメータの設定
# 辞書型で定義してsave_paramsに渡すとhparams.yamlが読みやすい？？
# argparseを使いこなせてないので、こんな書き方してます。すみません。
hp = {}
hp["11.lr"]             = 0.01
hp["12.batch_size"]     = 2
hp["21.max_epochs"]     = 5
hp["20.gpus"]           = 0
hp["51.train_data_x"]   = X
hp["51.train_data_y"]   = Y

# 作成したlightning modelと lightning moduleを生成
pl_model = LitModel(hp)
pl_module = LitModule(hp["51.train_data_x"], hp["51.train_data_y"], batch_size=hp["12.batch_size"])

# callbackの定義　earlystoppingとcheckpoint書き出しが良く使う設定
early_stopping = pl.callbacks.EarlyStopping("valid_loss")
model_chkpoint = pl.callbacks.ModelCheckpoint(save_last=True, save_top_k=2, monitor="valid_loss")
# trainerの生成 callbacksはリスト型で渡す
trainer = pl.Trainer(gpus=hp["20.gpus"], max_epochs=hp["21.max_epochs"], callbacks=[early_stopping, model_chkpoint])
# model, modeuleを渡してfitで学習実行
trainer.fit(pl_model, pl_module)

上記を実行すると、GPUの利用状況とか、モデルの緒言とかを書き出して、
プログレスバーを表示しながら学習を進めてくれる。
（今回の題材だと1epochが早すぎてプログレスバーが邪魔なだけだが…通常はありがたい）

GPU available: False, used: False
TPU available: False, using: 0 TPU cores

  | Name    | Type    | Params
------------------------------------
0 | model   | Model   | 2
1 | loss_fn | MSELoss | 0
------------------------------------
2         Trainable params
0         Non-trainable params
2         Total params
0.000     Total estimated model params size (MB)
Epoch 261: 100%|██████████████████████| 9/9 [00:00<00:00, 101.10it/s, loss=0.00124, v_num=14, valid_loss_epoch=0.000878, train_loss_step=0.00269, train_loss_epoch=0.00116, valid_loss_step=0.000486] 

参考文献

• ゼロから作るDeep Learning
  Deep Learningの基本を知るにはやっぱりこの本が分かりやすい

• PyTorch実践入門
  Pytorchで線形回帰する部分はこの本で勉強させていただきました。
  第1部はpytorch初学者にお勧めできると思います。第2部は…いつか読みます。

免責

　　コピペ自由ですが、責任は負えませんのでご了承ください。