機械学習とかの実験のパラメータをいい感じにやるコード

機械学習とかやっていると複数のパラメータを,コマンドラインから設定したいことがある. 割といい感じにできるようになった気がするので,備忘録として書いておく.

コマンドラインパーサとしてclickを使う. clickについては,こちらが詳しくていい感じだった.

この記事では以下ができるようにする.

  1. コマンドラインからいくつかのパラメータを設定する.
  2. パラメータを増やしたくなったときに,さくっとできるようにする.
  3. パラメータを名前に含んだディレクトリを作成して,その中に色々入れる.
import click
import datetime
import os
import pickle


def train(batch_size, learning_rate, nb_epochs, **kwargs):
    return 0, 0


def test(test_mode, **kwargs):
    print(test_mode)
    return 0


@click.command()
@click.option("-b", "--batch_size", type=int, default=12, help="mini batch size")
@click.option("-l", "--learning_rate", type=float, default=0.01)
@click.option("-n", "--nb_epochs", type=int, default=10, help="number of epochs")
@click.option("-m", "--test_mode", type=click.Choice(["hoge", "hogehoge"]))
@click.option("-o", "output_prefix", type=str, default="../results")
def main(**kwargs):
    net, loss_hist = train(**kwargs)
    test(**kwargs)

    output_prefix = kwargs.pop("output_prefix")
    output_dir = output_prefix + "/hoge_"
    now = datetime.datetime.now().strftime("%y%m%d%H%M%S")
    output_dir += now
    kwargs["output_dir"] = output_dir

    os.mkdir(output_dir)
    with open(output_dir + "/log.txt", "w") as f:
        for key, param in kwargs.items():
            print("{}, {}".format(key, param), file=f)
    with open(output_dir + "/params.pkl", "wb") as f:
        pickle.dump(kwargs, f)
    with open(output_dir + "/loss_hist.pkl", "wb") as f:
        pickle.dump(loss_hist, f)

    return


if __name__ == '__main__':
    main()

kwargsには辞書型が入ってくるので,あとはお好みで色々できる. 引数の最後に kwargsをdef train(batch_size, learning_rate, nb_epochs, **kwargs):な感じで入れておけば,必要な変数だけとってきて,残りはおいておくことができるので,いい感じに書けた.