Pytorchで学習コードを書くときの階層構造

## Pytorchで学習コードを書くときの階層構造

- モデルと訓練データ、バリデーションデータはfoldごとに作る。
- OOFの予測を出したいので、バリデーション時の予測は保持する。
- idxは`sklearn.model_selection`または`all_df`から作る。
- バリデーションでのメトリクスの評価は、バッチごとにできない場合があるので、なるべくバリデーション全体の`y_valid`で計算する。


```python
all_df = read_data()
y_oof = np.zeros(..., dtype=np.float32)

cv = create_cv(n_folds, all_df)

for fold in range(1, n_folds+1):
    model = create_model()
    model.to(DEVICE)

    output_dir = create_output_dir(fold)

    train_idx, valid_idx = create_idx(fold, cv)
    train_df = all_df.iloc[train_idx]
    valid_df = all_df.iloc[valid_idx]
    y_valid = valid_df[TARGET_COL].values

    y_fold = run_a_fold(
        model,
        train_df,
        valid_df,
        y_valid,
        output_dir
    )
    y_oof[idx_val] = y_fold
```

- 1foldごとにオプティマイザ、データセット、データローダーは作り直す。
- 最良のエポックの結果を保持する。

```python
def run_a_fold(
    model,
    train_df,
    valid_df,
    y_valid,
    output_dir
):
    
    optimizer, scheduler = create_optimizer()

    train_dataset = ...
    valid_dataset = ...

    train_loader = ...
    valid_loader = ...
    
    es = EarlyStopping()
    mc = ModelCheckPoint(model, output_dir)

    for epoch in range(1, n_epochs + 1):
        train_score = train_a_epoch(model, optimizer, train_loader)
        valid_score, y_valid_pred = valid_a_epoch(model, valid_loader, y_valid)

        if scheduler:
            scheduler.step()
        
        mc(valid_score)
        if mc.has_saved:
            y_valid_best_pred = y_valid_pred

        es(valid_score)
        if es.judge():
            break

    return y_valid_best_pred


class ModelCheckPoint:
    def __init__(self, model, output_dir):
        self.model = model
        self.output_dir = output_dir
        self.best_score = None
        self.has_saved = False

    def __call__(self, score)
        if self.best_score is None or score > self.best_score:
            self.save()
            self.best_score = score
            self.has_saved = True
        else:
            self.has_saved = False

    def save(self):
        os.makedirs(self.output_dir, exist_ok=True)
        torch.save(
            self.model.state_dict(), 
            os.path.join(self.output_dir, 'model_best.pth')
        )
    

class EarlyStopping:
    def __init__(self, patience=10):
       self.patience = patience
       self._initialize()
       
    def _initialize(self):
        self.best_score = None
        self._reset()

    def _reset(self):
        self.counter = 0

    def _increment(self):
        self.counter += 1

    def __call__(self, score):
        if self.best_score is None:
            self.best_score = score
        else:
            if score < self.best_score:
                self._increment()
            else:
                self._reset()
                self.best_score = score

    def judge(self):
        return self.counter >= self.patience
```

- `predict`はテストデータでも使用可能。

```python
def train_a_epoch(model, optimizer, train_loader):
    model.train()
    criterion = create_criterion()

    losses = []

    for xx, yy in train_loader:
        xx = xx.to(DEVICE)
        yy = yy.to(DEVICE)

        optimizer.zero_grad()
        output = model(xx)
        loss = criterion(output, yy)
        loss.backward()
        optimizer.step()
        
        losses.append(loss.item())

    average_loss = average(losses)
    return average_loss


def valid_a_epoch(model, valid_loader, y_valid):
    y_valid_pred = predict(model, loader)
    # 評価指標によってはエポックごとに計算して和を取ることができないものがある ex)AUC
    valid_score = calculate_score(y_valid_pred, y_valid)
    return valid_score, y_valid_pred


def predict(model, loader):
    model.eval()
    predicts = []
    for xx, _ in loader:
        xx = xx.to(DEVICE)

        with torch.no_grad():
            output = model(xx)

        predicts.append(output.cpu().detach().numpy())
    predicts = np.array(predicts)
    return predicts
```

- 各foldのモデルで推論し、まとめる。

```python
test_df = read_test_data()

test_dataset = ...
test_loader = ...

test_predictions = []

for fold in range(1, n_folds+1):
    modelpath = os.path.join(create_output_dir(fold), "model_best.pth")
    model = create_model(modelpath)
    model.to(DEVICE)

    y_test_pred = predict(model, test_loader)
    test_predictions.append(y_test_pred)

test_predictions = np.array(test_predictions).mean(axis=0)
```