HELLO CYBERNETICS

深層学習、機械学習、強化学習、信号処理、制御工学、量子計算などをテーマに扱っていきます

torch.jit を使ってみたのでメモ

 

 

follow us in feedly

はじめに

これはただのメモです。 もっと効率的なコンパイルのしどころがある気がしますし、そうであってほしいのですが、現状これしか分からなかったので書きます。注意点もメモしておきます。

Python on CPU

まずはPythonのコードとして実行します。デバイスはGoogle colab上のCPUです。

典型的なCNNで試していきます。PyTorchにはReshapeするだけの変換クラスは無いので、畳み込み層と線形層を分離して書いて、fowardreshape処理を挟む必要があります。テンソルのサイズも常に意識する必要があるので面倒くさいですが、意図せぬ形状になってわけの分からんネットワークを組んでしまうということが無いので、これはこれで丁寧に書く意識がついて良いかもしれません。

本題から逸れましたがコードは下記のようになります。

class MyConv(torch.nn.Module):
    def __init__(self):
        super(MyConv, self).__init__()
        
        self.convnet = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
        )
        
        self.linearnet = nn.Sequential(
            nn.Linear(64*28*28, 1024),
            nn.ReLU(),
            nn.Linear(1024, 10),
            nn.LogSoftmax(dim=1),
        )
        
    def forward(self, x):
        batch = x.size(0)
        x = self.convnet(x)
        x = x.reshape(batch, -1)
        return x

特にミニバッチ学習などは行いません。ネットワークの推論処理のスピードだけを比較します。したがってデータローダーは使わずに、単なる架空のテンソルデータを訓練データに使います。

X = torch.randn(128, 3, 28, 28, requires_grad=True)
y = torch.randint(0, 10, size=(128,))

入力データは requires_grad=Trueにしておいてラベルの方は False にしないと後でエラーが出ます。

モデルをインスタンス化して、Negative Log Likelihood Lossを利用します。負の対数尤度って、確率モデルの最尤推定を行うときの損失関数として一般的なものであって、別にクラス分類のロスを指すわけではないのですが、PyTorchではクラス分類に使うロス関数として使われています(普通にCrossEntropyとかにして欲しい)。

model_python = MyConv()
loss_fn = torch.nn.NLLLoss()
optimizer = torch.optim.Adam(params=model_python.parameters())

続いて学習の関数は下記です。loss 計算後は各オブジェクトが勝手にメッセージ送り合うのでメソッド呼ぶだけで勝手に勾配が算出され、パラメータが更新されます。値を取り出すときは、requires_grad=True となっているテンソルに関しては detach() で計算グラフを切ってから numpy()numpy.arrayにします。ココも地味に注意が必要です。

def train():
    y_pre = model_python(X)
    loss = loss_fn(y_pre, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    print("loss: ", loss.detach().numpy())

今回はJupyterを使ったのでマジックコマンドで時間を測りました。

%%time
for _ in range(20):
    train()

12sくらいでした。

Python on GPU

GPUでの処理を行う場合は、モデルのパラメータと入力出力のテンソルをそれぞれGPUに転送する必要があります。 optimizer もインスタンス化されるときにモデルのパラメータの参照を得るので、ちゃんとGPUに送った後のパラメータを渡してあげてください(下記の順番)。モデルをインスタンス化してオプティマイザーにパラメータを渡してから、モデルをGPUに送った場合は、オプティマイザーが正しく動作できません。

model_python.to("cuda")
optimizer = torch.optim.Adam(params=model_python.parameters())

続いて訓練コード

def train_gpu():
    y_pre = model_python(X.to("cuda"))
    loss = loss_fn(y_pre, y.to("cuda"))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    print("loss: ", loss.to("cpu").detach().numpy())

入力とラベルをGPUに転送することと、Pythonで表示する場合にCPUに戻してくることくらいが注意点です。

550msになりました。メチャクチャ高速ですね…!さすが。

Torch on CPU

さて、torch.jit を使ってモデルをTorchにコンパイルしてから使います。まずモデルのクラスを書くときに、torch.nn.Moduleを継承するのではなく、torch.jit.ScriptModuleを継承します。モデル作成の時点から JITを利用するか否かを決めて置かなければならないのは、なかなか厳しい制約です…。

また forward メソッドの方も、@torch.jit.script_method を利用する必要があります。できればPythonでのモデルとシームレスに使えるのが理想ですが…。致し方ないでしょう。

更に、モデルの層をインスタンス化するときにもtorch.jit.tracecallメソッドを呼び出してどのような計算処理がなされるのかを一旦トレースしておかなければならないようです。これは普通のPython関数をコンパイルするときも同様で、 torch.jit.trace(fn, example_input) の形式を取ります。コンパイラがこの計算処理の過程を静的グラフに変換しているようです。

※したがって、仮に、if 文処理があるようなPython関数をトレースした場合、トレースしたときの分岐処理が計算グラフとして構築され、他方の分岐は一切考慮されなくなることに注意してください。torch.jit.trace には if 文それ自体を計算グラフにする機能ありません。

class MyConv(torch.jit.ScriptModule):
    def __init__(self):
        super(MyConv, self).__init__()
        
        self.convnet = torch.jit.trace(nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
        ), torch.randn(1, 3, 28, 28))
        
        self.linearnet = torch.jit.trace(nn.Sequential(
            nn.Linear(64*28*28, 1024),
            nn.ReLU(),
            nn.Linear(1024, 10),
            nn.LogSoftmax(dim=1),
        ), torch.randn(1, 64*28*28))
    
    @torch.jit.script_method
    def forward(self, x):
        batch = x.size(0)
        x = self.convnet(x)
        x = x.reshape(batch, -1)
        return x

話は若干逸れましたが、ネットワークを書くときは上記のような感じなる気がします…。 もしかしたら、torch.jit.traceは各々の層に都度使う必要があるのだろうか…?(動いたからまあいいや)

あとは普通に学習のコードを書くだけなので省略

なぜかコイツも12s程かかりました。

Torch on GPU

さて、本丸のGPUで静的グラフを使うケースです。 結論から先にいうと、全く高速化はされていませんでした。。。 なんで!!!何か間違ってんの!?!!?!?!

まずモデルを書きます。先程のTorch on CPU のモデルをGPUに送れば良いかと思ったら、torch.jit.ScriptModuleto(device)メソッドを持っていないというバグみたいな仕様なので、モデルは書き直しです。

つまり、GPUを使うのかCPUを使うのか、更にはJITを利用するのかも予め決めてコードを書かなければならないという、CPUとGPUがおおよそシームレスだった生TensorFlowよりも扱いづらい素晴らしい機能となっています。正直何かの間違いであってほしいので、詳しい方指摘ください…。

ちなみに saveメソッドでセーブした後 load してくるときに、既に出来上がった計算グラフをCPUに配置するのかGPUに配置するのかは指定できるので、とりあえずCPU仕様で書いておいて、一旦セーブしてロードするときに配置場所を選ぶという書き方にしても良いかもしれないです。

class MyConv(torch.jit.ScriptModule):
    def __init__(self):
        super(MyConv, self).__init__()
        
        self.convnet = torch.jit.trace(nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
        ).to("cuda"), torch.randn(1, 3, 28, 28, device="cuda"))
        
        self.linearnet = torch.jit.trace(nn.Sequential(
            nn.Linear(64*28*28, 1024),
            nn.ReLU(),
            nn.Linear(1024, 10),
            nn.LogSoftmax(dim=1),
        ).to("cuda"), torch.randn(1, 64*28*28, device="cuda"))
    
    @torch.jit.script_method
    def forward(self, x):
        batch = x.size(0)
        x = self.convnet(x)
        x = x.reshape(batch, -1)
        return x
    
model_torch = MyConv()
optimizer = torch.optim.Adam(params=model_torch.parameters())

とりあえずコードは上記のようになり、書く層をコンストラクタでGPUに送り、更に torch.jit.trace でトレースするときに流すテンソルもしっかりGPUに配置しておきましょう。

まあ面倒くさいと言っても、やることは単純です(単純であるがゆえにわざわざやらずに隠蔽してほしかった)。 まあ、インスタンス生成時に層の処理をトレースしているとなれば、当然その時点で計算グラフの静的に配置して決めているわけですから、後からto.(device) できないのは分からないでも無いですね…?

あとは学習コードはGPUの方と同じです。 早くはなりませんでした。

おそらく学習コードの方それ自体も計算グラフにしてしまいたいところなのですが、現状、@torch.jit.script しても怒られるので、よく分かりません…。やはりココらへんはTensorFlow2.0に圧倒的な軍配が上がります。