PyTorchとTensorBoardでOptunaログの可視化

はじめに

Deep Learningのネットワーク開発では、可視化にmatplotlibを使うことが多いと思いますが、TensorBoardも有用です。TensorFlowを使う場合は可視化手段としてTensorBoardを使えば良いのですが、PyTorchの場合はどうすれば良いのでしょうか？これまではtensorboardXというPyTorchからTensorBoardを使えるようにしたライブラリを使う人が多かったと思いますが、PyTorch v1.1.0からTensorBoardがオフィシャルでサポートされるようになったので、PyTorchでもTensorBoardを使えば良いのです！公式サポートは、やはりありがたいですね。
またDeep Learningの性能向上においては、泥臭いハイパーパラメータチューニングが欠かせませんが、Preferred Networksが開発しOSSとして公開しているOptunaというハイパーパラメータ自動最適化フレームワークを使うと簡単にチューニングができます。
TensorBoardとOptunaを組み合わせて使えば、最適化のログも見やすく、管理しやすくなると思い実装してみました。コードはGithubにあげてあります。

可視化結果とコードの説明

この実装ではCIFAR-10でのクラス分類タスクを行っており、Network構造はPyTorchのチュートリアルにあった非常に簡単なものを使っているので、 あまり性能はでていません。性能うんぬんよりも、TensorBoardやOptunaの使い方の参考にしていただければ。

可視化結果

Optuna log visualization

以下のようにOptunaで実行したハイパーパラメータチューニングのログをTensorBoardで可視化することができます。左側のHyperparametersのチェックボックスをOn/Offすることで、表示するパラメータを選択することができます。

Optuna log visualization

Network graph visualization

Network構造も可視化できます。 PyTorch v1.3ではadd_graph()してTensorBoardのGRAPHSタブを見に行っても何も表示されなかったので、v1.4.0を使うようにしたところ表示されるようになりました。詳しくはこちらを参照してください。

Network graph visualization

Loss, accuracy results

add_scalars()でAccuracyとLossをグラフ化しています。tagを見てもらえれば何のグラフかは分かると思いますが、上段最左のグラフ①はOptuna最適パラメータでのtestデータの全クラス平均のAccuracyです。その右のグラフ②はclassごとのAccuracyで、x軸はclass番号になっており、10クラス分類なので0−９です。さらにその右のグラフ③はvalidデータのエポックごとのAccuracyでOptunaのtrial number分、プロットがあります。今回はエポック数は100、trial numberは50にしています。さらにその右のグラフ、つまり上段最右のグラフ④はOptunaのtrialごとのvalidデータの全クラス平均のAccuracyで、x軸はtrial numberになっています。中段のグラフ⑤はvalidデータのクラスごとのAccuracyで、x軸はclass番号です。最下段のグラフ⑥、⑦はそれぞれtrain、validデータのLossで共にx軸はエポック数です。

Loss, accuracy results

Images visualization

train、valid、testデータの画像も表示させました。今回はbatch sizeは32とし、batch size分を表示しています。

Images visualization

コードの説明

コードの全体は説明しませんが、掻い摘んで簡単にコードの説明をします。

    if args.optuna:
        # Hyperparameter tuning by using optuna
        study = optuna.create_study(direction='maximize')
        study.optimize(objective_variable(trainloader, validloader, writer), n_trials=args.optuna_trialnum)
        print('Best params : {}'.format(study.best_params))
        print('Best value  : {}'.format(study.best_value))
        print('Best trial  : {}'.format(study.best_trial))

        df = study.trials_dataframe()
        print(df)

        if args.tensorboard:
            df_records = df.to_dict(orient='records')

            for i in range(len(df_records)):
                df_records[i]['datetime_start'] = str(df_records[i]['datetime_start'])
                df_records[i]['datetime_complete'] = str(df_records[i]['datetime_complete'])
                value = df_records[i].pop('value')
                value_dict = {'value': value}
                writer.add_hparams(df_records[i], value_dict)

study = optuna.create_study(direction='maximize')でdirection='maximize'はobjectiveの戻り値を最大化するように最適化するという意味です。今回はvalid accuracyを戻り値にしているので最大化にしていますが、valid lossを戻り値にするのならばdirection='minimize'として最小化にする必要があります。ちなみにデフォルトはdirection='minimize'です。
df = study.trials_dataframe()でoptunaの結果を取得できます。これはpandasのDataFrameになっているので、to_dict()で辞書型に変換しています。datetime_sartとdatetime_completeはそのままadd_hparams()とするとエラーになるのでstringに変換しています。また、add_hparams()は引数にhparam_dictとmetric_dictを取るので、df_recordsからvalueをpopしてvalue_dictにしてadd_hparams()に渡しています。

def get_optimizer(trial, model):
    # Search Adam and SGD
    optimizer_names = ['Adam', 'SGD']
    optimizer_name = trial.suggest_categorical('optimizer', optimizer_names)

    weight_decay = trial.suggest_loguniform('weight_decay', 1e-10, 1e-3)
    lr = trial.suggest_loguniform('lr', 1e-5, 1e-1)

    if optimizer_name == optimizer_names[0]:
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    else:
        optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=weight_decay, momentum=0.9)
    return optimizer


def objective_variable(trainloader, validloader, writer):
    def objective(trial):
        global trial_num
        trial_num += 1

        model = network.Net().to(device)
        optimizer = get_optimizer(trial, model)
        criterion = nn.CrossEntropyLoss()

        # Training
        valid_accuracy = train(model, trainloader, validloader, optimizer, criterion, writer, trial_num)

        # Hyperparameter tuning will be done as return become max, since this code use direction='maximize'.
        return valid_accuracy

    return objective

今回、Optunaで最適化しているハイパーパラメータは、optimizer、learning rateとweight decayです。get_optimizer()という関数にしました。Optunaのチュートリアルを見ると、objective(trial)関数で目的関数を定義してstudy.optimize(objective, n_trials)で最適化を実行していますが、objective(trial)関数には引数を追加できないので、objective_variable()という高階関数を定義してobjective(trial)関数をreturnするようにしています。この書き方はこちらを参考にしました。

全体コード

https://github.com/livlea/pytorch_tensorboard_optuna/tree/v_0.1.2

参考サイト

https://github.com/pytorch/pytorch/releases/tag/v1.1.0
https://tech.515hikaru.net/2019-06-26-optuna-have-arg/