機械学習とかやっていると複数のパラメータを,コマンドラインから設定したいことがある. 割といい感じにできるようになった気がするので,備忘録として書いておく.
コマンドラインパーサとしてclickを使う. clickについては,こちらが詳しくていい感じだった.
この記事では以下ができるようにする.
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):
な感じで入れておけば,必要な変数だけとってきて,残りはおいておくことができるので,いい感じに書けた.