【Python】処理毎の所要時間が明確になるスマートなwith文 ＋ @contextmanagerの使い方

処理の所要時間が明確なコードはありがたい。

機械学習のロジックの書き方は人それぞれですが、
個人的に機械学習のコードを見てて「イケてる...」と思う書き方は「処理の所要時間」が明確なコードです。
やはり機械学習は処理に時間がかかるので、時間に関する情報が明確なコードは配慮が行き届いているなと感じます。
処理の所要時間はwith文とcontextmanagerを用いて、簡潔に出力することができます。
以下のテンプレートをご覧ください。

コードテンプレート

script.py

import time
from contextlib import contextmanager
from time import sleep

@contextmanager
def timer(title):
    t0 = time.time()
    yield
    print("{} - done in {:.0f}s".format(title, time.time() - t0))

def data_import():
    sleep(10)

def data_preprocessing():
    sleep(20)

def train():
    sleep(30)

def main():
    with timer("Process data_import"):
        data_import()
    with timer("Process data_preprocessing"):
        data_preprocessing()
    with timer("Process train"):
        train()


if __name__ == "__main__":
    with timer("Full model run"):
        main()

実行結果

Process data_import - done in 10s
Process data_preprocessing - done in 20s
Process train - done in 30s
Full model run - done in 60s

それぞれのメソッドでの処理時間、
そして、最後に処理全体の所要時間が出力されてます。
以上、ぜひご活用ください！