pytorchを使うときの小さな高速化

はじめに

pytorchなどを使ってモデルの学習をする際は場合によっては学習に数時間かかったり，何だったら数日かかったりすることがザラにあります．そういう場合は実行時間の定数倍の影響が大きくなってきます．例えば，計算量の定数倍が2から1に変わるだけでも，2日かかっていた学習が1日に短縮できたりします．競技プログラミングとかだと，主に計算量のオーダーの勝負になるのでそういった点では雰囲気が違いますね．今回はよく使える事例での簡単なを紹介します．

コード例

np.arrayのリストをTensor型に変換するコードはよく見ると思います．今回はこういう場面で使える高速化になります．

import time
import numpy as np
import torch

a = [np.ones([100]) for _ in range(100)]

# 8.44s
start = time.time()
for _ in range(10000):
    b = torch.tensor(a)
print(time.time()-start)

# 0.44s
start = time.time()
for _ in range(10000):
    b = torch.tensor(np.array(a))
print(time.time()-start)

# 0,32s
start = time.time()
for _ in range(10000):
    b = torch.from_numpy(np.array(a))
print(time.time()-start)

上の3つのコードの違いはほとんどないですが，実行時間では大きな差があります．どういう理由かは知らないですが，np.arrayのlistをTensor型に変換する場合は，直接torch.tensor()を使って変換するよりも，一旦np.array()を使って全体をnp.arrayに変換した実行時間が短くなります．

これは大きめのlistをTensor型に変換する場面が多い画像関連の学習の際には，割と目に見えて差が実感できると思います．ちなみにtorch.as_tensorを使った場合でもtorch.tensorを使って場合と同様の処理ができ，単体で使うと少し遅いです．

個人的には以下のようなコードで使っています

def convert2tensor(x, device):
    return torch.tensor(np.array(x, dtype=np.float32), device=device)

おわりに

人生は有限なのでモデルの学習に時間を溶かしすぎないようにしましょう．