Colabでレジュームしながらループを回す

最近は深層学習の新しい手法(そんなめちゃくちゃ新しいわけでもないんだけど個人的な興味にdrivenされてるところが大きい)を書いています。これは生まれて初めてarXivにも投稿するつもりです。

やはり新しい手法を書く場合、ハイパーパラメータによるその性能の変化を調べるためにもかなり網羅的な解析が必要になる。既存手法と比較して狂ったパフォーマンスが出るようなもんでもないわけで、その辺の誠実さは持っておきたい。

最近の僕は基本的に全てをColabで済ますのだが、さすがに網羅的な探索を行おうとするとColab Pro+だろうが継続時間が全く足りない。

たとえば、下のコードのみたいに(非常に愚かなループではあるが)提案モデルを含む3つのアーキテクチャーを作って、最適化、再起入力に使うレイヤー、バッチサイズ、アクティベーション関数などを総当たりで回すことを考える。MyModelはkerasのModelから作ったやつ。ちなみにepochは1000ぐらいにしておくけど、validationの損失関数をもとにEarly Stoppingがかかるようになっているのでどうでもいい。いまのアーキテクチャとデータだと平均的には300ぐらいで止まる。

project = 'newModel'
X_train, Y_train, X_test, Y_test = generate_data()
for m in [1, 2, 3]: # Model
    for opt in ['sgd', 'adam', 'Adadelta', 'adagrad', 'rmsprop']: # Optimizer
        for rec in ['rnn', 'lstm', 'bilstm']: # recurrent cell
            for bs in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]: # Batchsize
                for act in ['tanh', 'relu']: # Activation
                    for nHid in [1, 2, 3, 4, 5]: # Number of Hidden Layers
                        for nNeur in [1, 2, 5, 10, 20, 50, 100]: # Number of Neurons
                            model = MyModel(m, opt, rec, bs, act, nHid, nNeur, epoch=1000)
                            hist  = model.train(X_train, Y_train)
                            save_results(project, hist)                        

もっと効率的に探索できることはわかっています。一旦網羅的にデータとって色々と性質を知りたいんです。

そして一周回るごとにsave_results関数が走る。これの中身のイメージは

def save_results(project, hist):
    f = '/content/drive/MyDrive/'+project+'.pkl'
    res = pd.read_pickle(f)
    _ = pd.DataFrame(hist)
    res = pd.concat([res, _], axis = 0)
    res.to_pickle(f)

意味がわかってもらえればいいのでファイルパスの結合とかファイルの存在確認とかいろいろ省略してるけど、要はGoogle Drive上に結果を集約するpickleファイルを置いておいて、ループが1周終わるごとに新しい結果をつなげていくイメージ。

この形で網羅的な解析をやっているんだけど、まあ全然終わらん。かれこれ2週間ぐらい回し続けている。

回し続けているとは言ってもColabはPro+プランでも最長24時間で止まるので、これをうまく(ほぼ)回り続けるようにする必要がある。そのためには、現在ループがどこまで回ったかを知る必要がある。

つまり、いま現在ループが何周目なのかをカウントする変数が必要となるので、cntという形で作ってやる。cntは最初の一回のみ0に初期化されて、前回どこまでで終わったかをprev_cntとして読み出す。そのprev_cntはsave_resultsに値が保存されるときのpandas dfのインデックスをカウントにでもしてあげればいい。

project = 'newModel'
X_train, Y_train, X_test, Y_test = generate_data()
cnt = 0
prev_cnt = resume_loop(project)
for m in [1, 2, 3]: # Model
    for opt in ['sgd', 'adam', 'Adadelta', 'adagrad']: # Optimizer
        for bs in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]: # Batchsize
            for act in ['tanh', 'relu']: # Activation
                for nHid in [1, 2, 3, 4, 5]: # Number of Hidden Layers
                    for nNeur in [1, 2, 5, 10, 20, 50, 100]: # Number of Neurons
                        if cnt < prev_cnt:
                            cnt += 1
                            continue
                        model = MyModel(m, opt, bs, act, nHid, nNeur, epoch=1000)
                        hist  = model.train(X_train, Y_train)
                        save_results(project, hist, cnt)
                        cnt += 1

つまり、save_resultsでこんなふうにcntを書き出す設定にしておく。

def save_results(project, hist, cnt):
    f = '/content/drive/MyDrive/'+project+'.pkl'
    res = pd.read_pickle(f)
    _ = pd.DataFrame(hist, index=[cnt])
    res = pd.concat([res, _], axis = 0)
    res.to_pickle(f)

こうしておけば、

def resume_loop(project):
    f = '/content/drive/MyDrive/'+project+'.pkl'
    res = pd.read_pickle(f)
    prev_cnt = res.index.max()

とかで最後に保存された結果のcntを読み出せるので、あとはcntがprev_cntより小さい間だけcontinueしてあげればいい。

んで本当はここからcolabのスケジューラーで定期的に回せたら一番いいんだけど、スケジューラーで回そうとすると1日に1回の頻度の割にインスタンスの最大継続時間が4時間になってしまう。僕の場合はどうせ解析の進行状況は半日に1回ぐらいは見返すので、それを眺めてるときに気が向いたらランタイムを再起動してあげるとかすればいい。

んでColabのランタイムは最大5つまで同時に走るので、そこの並行処理がうまく走るように書き換えてあげる(たとえば結果の読み書きにあたって同じファイルを同時に開いたりしないように工夫してあげる必要がある。基本は別ファイルに書き出していって結果見るときだけファイル名のprefixとか頼りに統合したほうがいい)。これで(いくらColabが遅いとはいえ)多少はマシになる。本当は研究室の解析サーバーでぶん回したほうが早いんだけど、とにかく保守がめんどくさい。年間6万でこれなら別にいいです。