pytorch-pfn-extrasが便利という話
この投稿はrioyokotalab Advent Calendar 2020 7日目の投稿です。
実験ログをきれいに取りたい
PyTorchでコードで学習コードを書いた時の悩みの一つ。それは、ログの取り方だと思います。取りたい情報は損失、学習ステップ、学習率など、様々です。それぞれの情報が、取得するために様々な実装を必要とします。
今回は、pytorch-pfn-extrasのステマ紹介をします。
pytorch-pfn-extrasは実験のログをきれいに表示したり、ファイル出力してくれたりはもちろん。学習の進捗状況の表示や、Chainerではおなじみだった遅延評価モデルモジュールも提供してくれたりしています。
普段、自分が使わない昨日も多々あるので、全部紹介することはできないですが、わかる範囲で紹介します。
以下
import pytorch_pfn_extras as ppe
します。
Reporter
ppe.reporting.reportメソッドは記録したい値をppeモジュールに保存していきます。後々出てくるmanagerやextensionsで可視化、集計などを行えます。
ppe.reporting.report({'loss': loss.item()})
みたいな感じで使います。
学習管理
ppe.training.ExtensionsManagerクラスは学習epoch数や、現在の合計イテレーション数などを管理し、学習の停止タイミングを決めたりします。ppe.reporting.reportで吐き出した値を元に停止するタイミングを決定したりすることもできます。
また、with文で、
with manager.run_iteration():
以下にコードを書いてあげると、ppe.reporting.reportに出力する値は、LogReportで集計されるまで蓄積され、LogReportで出力するタイミングで平均によって集計されます。この機能は感覚的でない気がしていますが...。
また、学習ステップが進んでいることをppeに伝えることができ、学習イテレーションを一つ増やすことができたりします。
便利ツール
pytorch-pfn-extrasの学習便利ツール。それは、extensionsと呼ばれるものです。学習率の監視や、ログの出力、進捗バーの表示など、学習をわかりやすくするための諸々の機能が詰まっています。
EvaluatorLogReportMicroAveragePrintReportProgressBarParameterStatisticsPlotReportobserve_lrobserve_valuesnapshotVariableStatisticsPlot
これらはトリガーと呼ばれるものを指定することで、学習中の望むタイミングで動作させることができます。例えば、学習エポックの変わり目や、1000イテレーション経過時点、他にも、validation損失が最小になったタイミングなどなどです。
これ以外にも、必要に応じてExtensionクラスを自作したり、関数をExtension化したりすることができます。
これらの機能からいくつかピックアップします。
observe_lr
optimizerを渡すことで、optimizerの学習率を監視してくれます。監視している学習率はデフォルトではlrという値でppeモジュール内で管理されます。
LogReport
指定ステップ数ごとに学習中の様々な値をjson形式にファイル出力する。ログを出力する時点での、経過時間や学習エポック数、イテレーション数、ovserve_lrで監視している学習率、指定すれば、ppe.reporting.reportで出力した値などを出力してくれます。
snapshot
指定したタイミングでモデルのスナップショットを取ってくれます。トリガーには、検証損失が最小になったタイミングで起動するものもあるので、それを指定することで、ベストモデルだけを保存したりすることができます。
学習テンプレートコード
普段自分はこんな感じで書いています。学習コード自体のベースはこちらの記事を参照してください。
class Evaluator(ppe.training.extension.Extension): priority = ppe.training.extension.PRIORITY_WRITER def __init__(self, model, device, metrics, loader, prefix='valid/'): self.model = model self.device = device self.metrics = metrics self.prefix = prefix self.loader = loader def __call__(self, manager): self.model.eval() logs = {name: [] for name, metric in self.metrics.items()} with torch.no_grad(): for data in self.loader: data = to(data, device) out = self.model(**data) for name, metric in self.metrics.items(): met = metric(**out).item() logs[name].append(met) for name, value in logs.items(): ppe.reporting.report({ self.prefix + name: np.mean(value) }) manager = ppe.training.ExtensionsManager( model, optimizer, num_epochs, iters_per_epoch=len(train_loader), out_dir=out_dir ) standard_trigger = (1, 'epoch') manager.extend(E.observe_lr(optimizer=optimizer), trigger=standard_trigger) manager.extend(E.LogReport(trigger=standard_trigger)) manager.extend(E.PrintReport(['epoch', 'iteration', 'lr', 'train/loss', 'valid/loss', 'valid/acc', 'elapsed_time']), trigger=standard_trigger) manager.extend(Evaluator(model, device, {'loss': criterion, 'acc': metric}, valid_loader, prefix='valid/'), trigger=standard_trigger) manager.extend(E.snapshot(target=model, filename='best.pth'), trigger=ppe.training.triggers.MaxValueTrigger(key='valid/acc', trigger=standard_trigger)) manager.extend(E.snapshot(target=model, filename='model.pth'), trigger=standard_trigger) while not manager.stop_trigger: for data in train_loader: with manager.run_iteration(): model.train() optimizer.zero_grad() data = to(data, device) out = model(**data) loss = criterion(**out) ppe.reporting.report({ 'train/loss': loss.item() }) loss.backward() optimizer.step() scheduler.step()
実行すると
epoch iteration lr train/loss valid/loss valid/acc elapsed_time 1 534 7.61087e-05 1.0535 0.489182 0.825715 770.13 2 1068 9.96594e-05 0.471393 0.372857 0.872201 1538.18 3 1602 9.69772e-05 0.387294 0.367487 0.870103 2305.61 4 2136 9.17624e-05 0.34103 0.360847 0.879198 3073.83 5 2670 8.42963e-05 0.298277 0.363659 0.882929 3840.81 6 3204 7.49812e-05 0.257328 0.377409 0.881996 4608.64 7 3738 6.43194e-05 0.221985 0.373493 0.883085 5375.9 8 4272 5.28857e-05 0.189842 0.415019 0.874767 6144.7 9 4806 4.12964e-05 0.154903 0.431548 0.877332 6910.01 10 5340 3.01763e-05 0.126962 0.471258 0.875233 7677.17 11 5874 2.01249e-05 0.109875 0.49328 0.873134 8444.17
こんな感じにログを表示してくれて、out_dir/logにこのログと同じ内容のjsonファイルが作成されます。
1stepで学習する内容をすべてmanager.run_iteration()内部に置くことで、学習ステップが始まる前に起動して欲しいトリガーや学習ステップが終わった後に起動して欲しいトリガーなどの起動タイミングをコントロールすることができます。
Evaluatorは自作していますが、これは、aucなどの全推論結果を加味した指標を取りたい場合、ppe.training.extensions.Evaluatorを使うと不便だからです。
トリガーや、その他extensionsの使い方はgithubにあるdocumentationsを参照してください。
遅延評価モジュール
少し話が変わりますが、遅延評価モジュールの紹介をします。
ユースケース
EfficientNetという大きさをb0からb7までいろいろ変更できるモデルがあります。これは大きさを変更すると、最終層の出力ベクトルサイズも変わるのですが、これをコロコロ変えたい場合、最終層の全結合層の入力サイズを毎回変えなければなりません。また、EfficientNetからSE-ResNeXtに変える場合は？その度に最終層のベクトルサイズを調べ直すのは手間ですよね。しかもハードコーディングはかっこ悪い...。
Lazy Modules
class Model(nn.Module): def __init__(self, backbone): super().__init__() self.backbone = backbone self.pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Linear(2048, 5) def forward(self, x, **kwargs): out = self.backbone.extract_features(x) out = self.pool(out).squeeze(-1).squeeze(-1) out = self.fc(out) kwargs['out'] = out return kwargs
このように全結合層を書いているところを
class Model(nn.Module): def __init__(self, backbone): super().__init__() self.backbone = backbone self.pool = nn.AdaptiveAvgPool2d(1) self.fc = ppe.nn.LazyLinear(None, 5) def forward(self, x, **kwargs): out = self.backbone.extract_features(x) out = self.pool(out).squeeze(-1).squeeze(-1) out = self.fc(out) kwargs['out'] = out return kwargs
としてあげるだけで先ほどの悩みは半分ほど解決です。
あとは、遅延評価モデルなので、fcの入力サイズを評価する必要があります。それを行うためには、
model = Model(backbone)
for example in train_loader:
break
with torch.no_grad():
out = model(**example)
と実行してあげて、一度、モデルにデータを通します。そうすることで、モデルはデータが入力された際の全結合層への入力ベクトルサイズを確認することができ、あとは、通常のモデルとして扱うことができるようになります。
まとめ
他にも、Configの管理ツールやONNXをよしなにしてくれる機能もあるらしいですが、使ったことがないので、今回は取り扱いません。
pytorch-pfn-extrasは実験のログ管理ツールの中では、かなりコンパクトで取り回しやすいものだと感じています。
ドキュメント早く書いて...
