PyTorchLightning my poor Refactoring

# PyTorch -> PyTorchLightning

https://www.youtube.com/watch?v=DbESHcCoWbM

見たので、動画とdocumentationを見ながら、手元の image classification model を refactoring してみました。
PyTorchLightningを活用しきれていないです。
```python
%pip install pytorch-lightning
import pytorch_lightning as pl
print(pl.__version__)
# '1.7.1'（2022.8）
```

## Grasp the picture of my project
1. make CONFIG class
2. import libraries
3. define
4. make_directory, make_pathlist
5. transform
6. Dataset
7. PyTorch_Lightning
8. Trainer

## 1. make CONFIG class
```python
import torch
import pprint
import datetime
dt_now = datetime.datetime.now()

class CFG:
    setting = {"MODEL_NAME" : "efficientnet_b3", "model_library" : "torchvision", "pretrained" : True, "Version" : "trial"}
    seed = 0
    num_workers = 4
    n_class = 0 # 後で上書き
    resize = 224
    batch_size = 32
    epochs = 5
    model_lr = 1e-3
    weight_decay = 1e-4
    beta1 = 0.9
    beta2 = 0.999
    mean = (0.485, 0.456, 0.406)
    std = (0.229, 0.224, 0.225)
    device = (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
    n_fold = 5
    criterion = "nn.CrossEntropyLoss()"
    optimizer = "optim.Adam(self.parameters(), lr=CFG.model_lr)"
    scheduler = "lr_scheduler.CosineAnnealingLR(optimizer, T_max=CFG.epochs)"

print(dt_now)
print(f"MODEL_NAME : {CFG.setting['MODEL_NAME']}")
print(f"model_library : {CFG.setting['model_library']} / pretrained : {CFG.setting['pretrained']}")
print(f"Version : {CFG.setting['Version']}")
print(f"batch_size : {CFG.batch_size}")
print(f"optimizer : {CFG.optimizer}")
print(f"scheduler : {CFG.scheduler}")
```
2022-08-23 00:40:12.595750
MODEL_NAME : efficientnet_b3
model_library : torchvision / pretrained : True
Version : trial
batch_size : 32
optimizer : optim.Adam(self.parameters(), lr=CFG.model_lr)
scheduler : lr_scheduler.CosineAnnealingLR(optimizer, T_max=CFG.epochs)

optimizer とか scheduler とか変更する可能性があるものはここに文字列として入れてます。
```python
optimizer = eval(CFG.optimizer)
```
CFGのクラス内変数は文字列ですが、eval関数を使って、参照させてます。

## 2. import libraries
今回使ってないものも入ってます...闇鍋状態...
```python
import os
import glob
import sys
import pandas as pd
import numpy as np
import random

from PIL import Image
import matplotlib.pyplot as plt
import seaborn as sns
sns.set()
# sns.set_style("ticks")

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
import torch.optim as optim
from torch.optim import lr_scheduler

from torchvision import datasets, models, transforms
import timm

import pytorch_lightning as pl
from pytorch_lightning import loggers as pl_loggers
from pytorch_lightning import Trainer, seed_everything
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks import ModelCheckpoint, RichProgressBar
from pytorch_lightning.callbacks.progress.rich_progress import RichProgressBarTheme
from pytorch_lightning.callbacks.progress import ProgressBarBase

from sklearn.model_selection import train_test_split, StratifiedKFold

import torchmetrics
from torchmetrics import Accuracy, MetricCollection, Precision, Recall

import warnings
warnings.filterwarnings('ignore')
```

## 3. define
```python
def torch_seed(seed=0):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.use_deterministic_algorithms = True

def load_model():
    if CFG.setting["model_library"] == "torchvision":
    # torchvision
        net = eval(f'models.{CFG.setting["MODEL_NAME"]}')(pretrained=CFG.setting["pretrained"])
        
        # resnet
        if "resne" in CFG.setting["MODEL_NAME"]:
            fc_in_features = net.fc.in_features
            net.fc = nn.Linear(fc_in_features, CFG.n_class)
        # vgg
        if "vgg" in CFG.setting["MODEL_NAME"]:
            fc_in_features = net.classifier[6].in_features
            net.classifier[6] = nn.Linear(fc_in_features, CFG.n_class)
        # efficient
        if "efficientnet_" in CFG.setting["MODEL_NAME"]:
            fc_in_features = net.classifier[1].in_features
            net.classifier[1] = nn.Linear(fc_in_features, CFG.n_class)
        # vit
        if "vit_" in CFG.setting["MODEL_NAME"]:
            fc_in_features = net.heads.head.in_features
            net.heads.head = nn.Linear(fc_in_features, CFG.n_class)


    elif CFG.model_library == "timm":
        net = timm.create_model(CFG.MODEL_NAME, pretrained = CFG.pretrained, num_classes = CFG.n_class)

    return net

def visualize_process(train_loss_log, valid_loss_log, train_acc_log, valid_acc_log):
    """
    train_loss_log, valid_loss_log, train_acc_log, valid_acc_log : list
    """
    fig = plt.figure(figsize=(15,5))
    ax = plt.subplot(1, 2, 1)
    plt.plot(train_loss_log, label="train")
    plt.plot(valid_loss_log, label="valid")
    plt.title("Loss_log")
    plt.xlabel("epoch")
    plt.ylabel("Loss")
    plt.legend()

    ax = plt.subplot(1, 2, 2)
    plt.plot(train_acc_log, label="train")
    plt.plot(valid_acc_log, label="valid")
    plt.title("Acc_log")
    plt.xlabel("epoch")
    plt.ylabel("accuracy")
    plt.legend()

    plt.show()
    return fig
```
torch_seed : seed固定
load_model : CFGに基づいてmodelをloadします。
visualize_process : 訓練・検証ログのlistからグラフを描きます

## 4. make_directory , make_pathlist
```python
rootpath = "/root/workspace/Pytorch_Lightning/data"
```

```
rootpath
    ├── class_1
    │     ├── class_1の画像1.jpg
    │     ├── class_1の画像2.jpg
    │     └── class_1の画像n.jpg
    ├── class_2
    │     ├── class_2の画像1.jpg
    │     ├── class_2の画像2.jpg
    │     └── class_2の画像n.jpg
    ├── class_3
    │     ├── class_3の画像1.jpg
    │     ├── class_3の画像2.jpg
    │     └── class_3の画像n.jpg
    └── class_...
          ├── class_...の画像1.jpg
          ├── class_...の画像2.jpg
          └── class_...の画像n.jpg

```
ありがちなdierctoryを想定してます。
rootpath 以下のフォルダを数えて、分類クラスの数も確定させてます。

```python
# rootpath
rootpath = "/root/workspace/Pytorch_Lightning/data"

# param_save_path
parampath = f"/root/workspace/Pytorch_Lightning/params/{CFG.setting['MODEL_NAME']}_{CFG.setting['Version']}"
os.makedirs(parampath, exist_ok=True)

print(f"rootpath : {rootpath}")
print(f"parampath : {parampath}")

# listdirの下に .DS_Store が勝手に生成されてることがあるので対策してます...
class_name_list = [class_name for class_name in os.listdir(rootpath) if class_name != ".DS_Store"]
CFG.n_class = len(class_name_list) # クラスの数

filepath = os.path.join(rootpath , "**/*.jpg")
path_list = [path for path in glob.glob(filepath)]

print(f"n_path : {len(path_list)}")
print(f"n_class : {CFG.n_class}")
```
rootpath : /root/workspace/Pytorch_Lightning/data
parampath : /root/workspace/Pytorch_Lightning/params/efficientnet_b3_trial
n_path : 144
n_class : 3

**pathの指定は環境によって変更が必要です。**

## 5. transform
```python
class train_ImageTransform():
    def __init__(self, resize, mean, std):
        self.data_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((resize, resize)),
            transforms.RandomRotation(degrees=(-10, 10)), # rotate
            transforms.RandomHorizontalFlip(p=0.5), # Flip
            transforms.Normalize(mean, std), # Normalize
            transforms.RandomErasing(p=0.8, scale=(0.06, 0.20), ratio=(0.1, 4.0)), # mask
            ])
    def __call__(self, img):
        return self.data_transform(img)

class valid_ImageTransform():
    def __init__(self, resize, mean, std):
        self.data_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((resize, resize)),
            transforms.Normalize(mean, std),
            ])
    def __call__(self, img):
        return self.data_transform(img)
```

## 6. Dataset
画像ファイルへの path と label の2列からなる DataFrame を作って、そこから Dataset を作ることにします。

```python
# make_pathlist
filepath = os.path.join(rootpath , "**/*.jpg")
path_list = [path for path in glob.glob(filepath)]
class_name_list = [class_name for class_name in os.listdir(rootpath) if class_name != ".DS_Store"]
# make_class / class_name(str) -> label(int) の変換を行うdictを作ります
label_map = dict()
for i, class_name in enumerate(class_name_list):
    label_map[class_name] = i
# 作ったdictでlabel_listにします。intでリストに格納されます
label_list = [label_map[path.split("/")[-2]] for path in path_list]
# DataFrame にする
df = pd.DataFrame({
    "path" : path_list,
    "label" : label_list
})

# print(df.shape) # (len(file_list),2)
```
この df から Dataset を作るクラスを作ります。
```python
class ImageDataset(Dataset):
    def __init__(self, df, transform=None):
        super().__init__()
        self.df = df
        self.transform = transform 
    
    def __len__(self):
        return len(self.df) 

    def __getitem__(self, index):
        row = self.df.iloc[index]
        img_path = row.path
        img = self.transform(Image.open(img_path).convert("RGB"))
        label = row.label
        return img, label
```
動作確認
```python
train_transform = train_ImageTransform(CFG.resize, CFG.mean, CFG.std)
train_dataset = ImageDataset(df, transform=train_transform)
print(len(train_dataset)) # n_pathlist
print(train_dataset[0][0].shape, train_dataset[0][1]) # torch.Size([3, 224, 224]) labelを表すint
```
この df は StratifiedKFold して使うことになります。


## 7. PyTorch_Lightning
datasetができたら PyTorchLightning を使います。
- DataLoaderの作成
- 訓練・検証ループの作成
の2つを作ります。

train_datasetとvalid_datasetを渡すとこのクラスででDataLoaderが作れます。

```python
# make DataLoader for Pytorch_Lightning
class PlDataModule(pl.LightningDataModule):
    def __init__(self, train_dataset, valid_dataset):
        super().__init__()
        self.train_dataset = train_dataset
        self.valid_dataset = valid_dataset

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=CFG.batch_size, shuffle=True, num_workers=CFG.num_workers)

    def val_dataloader(self):
        return DataLoader(self.valid_dataset, batch_size=CFG.batch_size, shuffle=False, num_workers=CFG.num_workers)
```
PyTorchLightningでLoop部分を作ります。
```python
# Loop
class Net(pl.LightningModule):
    # make_model
    def __init__(self, cv):
        super(Net, self).__init__()
        # model
        self.model = load_model()
        # cv
        self.cv = cv
        # other
        self.lr = CFG.model_lr
        self.criterion  = eval(CFG.criterion)
        # metrics
        self.classes = CFG.n_class
        self.accuracy = torchmetrics.Accuracy()
        self.F1score = torchmetrics.F1Score(num_classes=self.classes, average='macro')
        # log_list
        self.train_loss_log, self.train_acc_log = list(), list()
        self.valid_loss_log, self.valid_acc_log = list(), list()
        self.F1_log = list()
    
    # forward
    def forward(self, x):
        x = self.model(x)
        return x

    # optimizer & scheduler
    def configure_optimizers(self):
        optimizer = eval(CFG.optimizer)
        scheduler = eval(CFG.scheduler)
        return [optimizer],[scheduler]
    
    # 1.training_step
    def training_step(self, batch, batch_idx):
        # x:image y:label
        imgs, labels = batch # imgs:[B, C, H, W], labels:[B]
        logits = self.forward(imgs) # [B, class]
        loss = self.criterion(logits, labels)
        # display log
        self.log("train_acc", self.accuracy(logits, labels),prog_bar=False,logger=True,
                            on_epoch=True, on_step=False)
        print(".", end='')
        return {"loss" : loss, "logits" : logits, 
                    "labels" : labels, "batch_loss" : loss.item()*imgs.size(0)}
        # returnのDict内の"loss"を基にモデルがアップデートされる

    # 2.validation_step
    def validation_step(self, batch, batch_idx):
        imgs, labels = batch
        logits = self.forward(imgs)
        loss = self.criterion(logits, labels)
        # display log
        self.log("valid_acc", self.accuracy(logits, labels),prog_bar=False,logger=True,
                            on_epoch=True, on_step=False)
        return {"valid_loss" : loss, "logits" : logits,
                    "labels" : labels, "batch_loss" : loss.item()*imgs.size(0)}

    # 3.validation_epoch_end
    def validation_epoch_end(self, val_step_outputs):
        logits = torch.cat([o["logits"] for o in val_step_outputs], dim=0) # [n_train_dataset, class]
        labels = torch.cat([o["labels"] for o in val_step_outputs], dim=0) # [n_train]
        epoch_loss = sum([o["batch_loss"] for o in val_step_outputs]) / logits.size(0)
        acc = self.accuracy(logits, labels)
        F1 = self.F1score(logits, labels)
        self.valid_loss_log.append(epoch_loss)
        self.valid_acc_log.append(acc.item())
        self.F1_log.append(F1.item())
        # display log
        self.log("valid_acc_epoch", acc, prog_bar=True,logger=True,
                            on_epoch=True, on_step=False)
        # # save記述
        if np.argmax(self.valid_acc_log) == self.current_epoch:
            torch.save(self.model.state_dict(), os.path.join(parampath, f"best_acc_{self.cv}.pth"))

    # 4.training_epoch_end
    def training_epoch_end(self, train_step_outputs):
        # リストの中に return の dict が入ってる
        # print(train_step_outputs) # [ [training_step_outputs], [training_step_outputs],...,[training_step_outputs] ]
        logits = torch.cat([o["logits"] for o in train_step_outputs], dim=0) # [n_train_dataset, class]
        labels = torch.cat([o["labels"] for o in train_step_outputs], dim=0) # [n_train]
        epoch_loss = sum([o["batch_loss"] for o in train_step_outputs]) / logits.size(0)
        acc = self.accuracy(logits, labels)
        self.train_loss_log.append(epoch_loss)
        self.train_acc_log.append(acc.item())
        # display log
        self.log("train_acc_epoch", acc, prog_bar=True,logger=True,
                            on_epoch=True, on_step=False)
        print()
        print(f"epoch : {self.current_epoch} >>> ", end='')
        print(f"train_loss : {self.train_loss_log[-1]:.4f} / train_acc : {self.train_acc_log[-1]:.4f} ", end='') 
        print(f"/ valid_loss : {self.valid_loss_log[-1]:.4f} / valid_acc : {self.valid_acc_log[-1]:.4f} / F1 : {self.F1_log[-1]:.4f}")
        print()
```
- __init__の部分でmodel, learning_rate, criterion, 評価指標の計算機とかを決めます。
- forwardは [B, C, H, W]の画像のテンソルを渡した時の挙動で、モデルに突っ込むよということです。
- configure_optimizers で optimizerとscheduler を設定
- training_step -> validation_step -> validation_epoch_end -> training_epoch_end という順番で動いていきます。記述の順番に制約はありません。
- 訓練回数は self.current_epoch で取得できます。
- self.logはlogが残せます。この辺りがまだあまり理解できていません。
progress_barを表示した時に載せるような設定にできます。
あとは tensor_board のような logger に送ることができたりという設定ができるようです。
あまり理解ができてない故、強引に表示してます。ここはもっとスマートにできると思われます...。
- PyTorch Lightningでは net.train(),net.eval(),to(device),optimizer.zero_grad()とかもろもろの記述が不要、というか書かないでも中で自動でやってくれているようです。

- CrossValidationで回したいので、cvという変数をそのまま追加しています。訓練中に検証時のacc最高のモデルを保存するようにしています。

- validation_epoch_endの引数のval_step_outputs の部分は
printすれば分かりますが、リストの中に validation_step の return が dict で入ってます。
epoch_endの中ではこの dict を処理する記述をしてます。

```python
import torch

d1 = {'loss': 1, 'labels': torch.tensor([0,1,2]), 'logits': torch.tensor([[1,1,1], [2,2,2], [3,3,3]])}
d2 = {'loss': 10, 'labels': torch.tensor([1,1,1]), 'logits': torch.tensor([[10,10,10], [20,20,20], [30,30,30]])}
d3 = {'loss': 100, 'labels': torch.tensor([2,2,0]), 'logits': torch.tensor([[0.1,0.1,0.1], [0.2,0.2,0.2], [0.3,0.3,0.3]])}
outputs = [d1, d2, d3]
print(outputs)

for o in outputs:
    print("*****")
    print(o)

print()
sum_ = sum([o["loss"] for o in outputs])
print(sum_)
cat_ = torch.cat([ o["logits"] for o in outputs ], dim=0)
print(cat_)
print(cat_.shape)
```
こんな感じのことをしてるはずです。

## 8. Trainer
PyTorchLightningのdocumentでは以下のような記述が例として記述されています。
```python
model = MyLightningModule()
trainer = Trainer()
trainer.fit(model, train_dataloader, val_dataloader)
```
今回は MyLightningModule() は Net() と記述、dataloaderは PlDataModuleクラスで作成、cvの回数をそのまま記述したので、
```python
data_module = PlDataModule(train_dataset, valid_dataset)
# Trainer
model = Net(cv=cv)
trainer = pl.Trainer(max_epochs=CFG.epochs, gpus=1, callbacks=progress, num_sanity_val_steps=0)
trainer.fit(model,datamodule=data_module)
```
のように書くことにします。
max_epochsは訓練回数。並列・分散しないのであれば、gpus=1。
callbacksの部分では progress_barを消す設定をしました。progress_barの制御がVSCodeで難しかったので...
```python
# ProgressBar setting
class NotebookProgress(ProgressBarBase):
    def __init__(self):
        super().__init__()
        self.enable = True
    def disable(self):
        self.enable = True

progress = NotebookProgress()
```
callbacksでは他にも様々な設定ができるようです。checkpointとかearlystoppingとか。

num_sanity_val_stepsの指定ですが、ここを指定しないと、training_step -> validation_step -> validation_epoch_end -> training_epoch_end のループに入る前に、訓練していない状態で validation が先に1回だけ回るという設定がデフォルトになっています。

https://pytorch-lightning.readthedocs.io/en/1.7.2/common/trainer.html#num-sanity-val-steps

こちらに設定の意図などとても分かりやすく書いてありました。

では最後に、StratifiedKFoldの枠組みで Trainer を回します。

```python
# fix seed
pl.seed_everything(CFG.seed)
# ProgressBar setting
progress = NotebookProgress()
# transform
train_transform = train_ImageTransform(CFG.resize, CFG.mean, CFG.std)
valid_transform = valid_ImageTransform(CFG.resize, CFG.mean, CFG.std)

skf = StratifiedKFold(CFG.n_fold, shuffle=True, random_state=CFG.seed)
cv = 0
cv_log = {}

for train_index, valid_index in skf.split(df, df["label"]):
    cv +=1
    print(f"************** fold: {cv} ***************")
    # df => train_df, valid_df
    train_df, valid_df = df.iloc[train_index], df.iloc[valid_index]
    # dataset
    train_dataset = ImageDataset(train_df, transform=train_transform)
    valid_dataset = ImageDataset(valid_df, transform=valid_transform)
    # data_module
    data_module = PlDataModule(train_dataset, valid_dataset)
    model = Net(cv=cv)
    trainer = pl.Trainer(max_epochs=CFG.epochs, gpus=1, callbacks=progress, num_sanity_val_steps=0)
    trainer.fit(model,datamodule=data_module)
    # plot
    fig = visualize_process(model.train_loss_log, model.valid_loss_log, model.train_acc_log, model.valid_acc_log)

    print(f"***** fold:{cv}:best_acc *****")
    print(max(model.valid_acc_log))
    print()
    # cvごとのlog
    # cv_logはmapになっていて、cvの数字によって log を保存
    cv_log[f"cv{cv}_train_loss_log"] = model.train_loss_log
    cv_log[f"cv{cv}_train_acc_log"] = model.train_acc_log
    cv_log[f"cv{cv}_valid_loss_log"] = model.valid_loss_log
    cv_log[f"cv{cv}_valid_acc_log"] = model.valid_acc_log
    cv_log[f"cv{cv}_valid_F1_log"] = model.F1_log

# 訓練終了後に一番良い値を表示する
for i in range(CFG.n_fold):
    print(f"********** cv{i+1} **********")
    best_acc = max(cv_log[f"cv{i+1}_valid_acc_log"])
    best_acc_index = np.argmax(cv_log[f"cv{i+1}_valid_acc_log"])
    print("best_acc_index: {}, best_acc: {:.04f}".format(best_acc_index, best_acc))
    best_loss = min(cv_log[f"cv{i+1}_valid_loss_log"])
    best_loss_index = np.argmin(cv_log[f"cv{i+1}_valid_loss_log"])
    print("best_loss_index: {}, best_loss: {:.04f}".format(best_loss_index, best_loss))
    best_F1 = max(cv_log[f"cv{i+1}_valid_F1_log"])
    best_F1_index = np.argmax(cv_log[f"cv{i+1}_valid_F1_log"])
    print("best_F1_index: {}, best_F1: {:.04f}".format(best_F1_index, best_F1))
```

## poorな部分
- 今回、訓練結果を強引に表示していますが、これはスマートなやり方がありそうです。
- CrossValidation 自体もこのように for 文で記述しなくても、どこかで定義できる方法があると思われます。

またdocumentを見て改善します...

## 参考
https://www.youtube.com/watch?v=DbESHcCoWbM&t=3s

https://pytorch-lightning.readthedocs.io/en/latest/