最小のpytorchを作る。

前回のブログで、PyTorchのlibtorchをC++言語から使用する方法を紹介しました。
このlibtorchはtorch.dll、c10.dll、caffe2.dllの３つのモジュールからできています。torch.dllが20Mバイト、caffe2.dllが30Mバイトあります。
うーんでかすぎる。
ソースコードを見る限り、どうみてもcaffe2.dllなんて使っていないように見る。

また、C++の部分はソースコードがビルドされたら自動でモジュールを登録するように実装されているため、全部をビルドしなくてもある程度動くようになっています。

これ、もっと小さくできるんじゃない？
もっと小さくしないと組み込みの製品やスマホアプリに使えないじゃないか。

そこで本日はMNISTが動く、最小のlibtorchをCmakeを使わず、必要なファイルだけを手動でVisualStudioに追加していき、手動でビルドしてみたいと思います。


まず、c10、Aten、cpuinfo、eigen、TH、THNNの各ディレクトリのソースを全部手動でVisaualStudioに追加します。
ついでに、初回ビルド時にツールによって生成される、C++のヘッダーファイルとソースコードをbuild/c10、build/atenのディレクトリから取ってきて追加します。


次に今回はdllを作るわけではないので、ソースコード中のdllimportやdllexportのdefine部分を何か所か以下のように修正します。

修正前:

#define TH_CPP_API __declspec(dllimport)
#define TH_CPP_API __declspec(dllimport)

修正後:

#define TH_CPP_API
#define TH_CPP_API

ファイル名にcudaやmiopenやhipがついているものはGPU用のコードなのでビルドしなくて大丈夫です。
cudaなどの名前がついていないソースをビルドをします。
そんなにファイル数が多くないので、簡単ですね。

一か所「caffe2/hogeohoge.hがインクルードできません。」と言われますがcaffe2のソースなんて使っていないのでコメントアウトします。
一部の関数がダブっているとコンパイラに怒られますが、意外とスムーズに全部のソースをビルドすることができます。


上記のものをビルドすると、リンク時に「torch::*が見つかりません。」と多量のリンクエラーがでます。
torch/csrcディレクトリの中のapi,jit,autogradをビルドしていきます。
ここに一つ問題が。
torchのディレクトリはPythyon用のインタフェースコードとC++API用のコードが混在しています。
なので、#include<Pyton.h>などがないものをビルドしていきます。
jitの中にある、export.cppとimport.cppはProtoBufを使うようになっています。サンプルのMNIST.cppなど、jitのファイル入出力を行わない場合は空のオブジェクトを返す何もしない関数に修正します。

torchディレクトリには初回ビルド時に自動生成される以下の巨大ファイルがあります。

VariableTypeEverything.cpp
VariableType_0.cpp
VariableType_1.cpp
VariableType_2.cpp
VariableType_3.cpp
VariableType_4.cpp

これらのファイルがなくてもビルドはできますが、実行時にエラーになるため、これらのファイルも忘れずにビルドをします。

こんな適当なビルドでほんとうに動くのか自分でも疑問でしたが、とりあえずビルドとおるじゃん。しかもあっさり、MNIST.cppが動きました。

なんだ、PyTorchのlibtorchってc10、Aten、cpuinfo、eigen、TH、THNNと、torch/csrcディレクトリのなかのapi、jit、autogradだけコンパイルをすれば動くんですね。
この部分だけビルドして公開してくれればいいのに。
これで組み込みとかスマホアプリでC++版のPyTorchが使えます。

Python版のPyTorchとC++版のPyTorchってどれくらい差があるのでしょう。
両方のMNISTのソースコードを比べてみました。

Python版
--------------------
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4*4*50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)
--------------------


C++版
--------------------
struct Net : torch::nn::Module {
    Net()
    : conv1(torch::nn::Conv2dOptions(1, 10, /*kernel_size=*/5)),
    conv2(torch::nn::Conv2dOptions(10, 20, /*kernel_size=*/5)),
    fc1(320, 50),
    fc2(50, 10) {
        register_module("conv1", conv1);
        register_module("conv2", conv2);
        register_module("conv2_drop", conv2_drop);
        register_module("fc1", fc1);
        register_module("fc2", fc2);
    }
    torch::Tensor forward(torch::Tensor x) {
        x = torch::relu(torch::max_pool2d(conv1->forward(x), 2));
        x = torch::relu(
        torch::max_pool2d(conv2_drop->forward(conv2->forward(x)), 2));
        x = x.view({-1, 320});
        x = torch::relu(fc1->forward(x));
        x = torch::dropout(x, /*p=*/0.5, /*training=*/is_training());
        x = fc2->forward(x);
        return torch::log_softmax(x, /*dim=*/1);
    }
    torch::nn::Conv2d conv1;
    torch::nn::Conv2d conv2;
    torch::nn::FeatureDropout conv2_drop;
    torch::nn::Linear fc1;
    torch::nn::Linear fc2;
};
--------------------

なんと、ほとんど同じですね。

Pythonで書いてGPU付きのPCで学習させた結果をどんどんC++で使えそうですね。

今回紹介した方法で、Python版のPytorchで作ったモジュールを、即C++でアプリケーションや製品に組み込むことができます。これは便利。

本日は組み込みやスマホアプリ用の最小のPyTorchを作る方法でした。