WindowsのPyTorchでCUDAを使ったカスタムレイヤーを作る。

Windows上のPyTorchで、CUDAを使ってカスタムレイヤーを作りたい、というニッチな人向けの情報です。

https://github.com/sniklaus/pytorch-extension に、PyTorchでCUDAを使ったカスタムレイヤーを作るためのサンプルがあるのですが、そのままだとWindowsでは動かなかったので、無理やり動かしました。

前提条件

clのバージョン

> cl
Microsoft(R) C/C++ Optimizing Compiler Version 19.11.25547 for x64

Visual Studio 2017 です。

NVCCのバージョン

> nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2017 NVIDIA Corporation
Built on Fri_Nov__3_21:08:12_Central_Daylight_Time_2017
Cuda compilation tools, release 9.1, V9.1.85

PyTorchのバージョン

> python -c "import torch; print(torch.__version__)"
0.3.1.post2

conda install -c peterjc123 pytorch cuda90でインストールしました。

クローン

> vcvarsall.bat x86_amd64
> git clone https://github.com/sniklaus/pytorch-extension
> cd pytorch-extension

ソース修正

install.pyの修正

diff --git a/install.py b/install.py
index 1a26280..508db3c 100644
--- a/install.py
+++ b/install.py
@@ -10,7 +10,7 @@ strObjects = []

 if torch.cuda.is_available() == True:
        strHeaders += ['src/HadamardProduct_cuda.h']
-       strSources += ['src/HadamardProduct_cuda.c']
+       strSources += ['src/HadamardProduct_cuda.cpp']
        strDefines += [('WITH_CUDA', None)]
        strObjects += ['src/HadamardProduct_kernel.o']
 # end
@@ -25,7 +25,9 @@ objectExtension = torch.utils.ffi.create_extension(
        relative_to=strBasepath,
        include_dirs=[os.path.expandvars('$CUDA_HOME') + '/include'],
        define_macros=strDefines,
-       extra_objects=[os.path.join(strBasepath, strObject) for strObject in strObjects]
+       extra_objects=[os.path.join(strBasepath, strObject) for strObject in strObjects],
+        library_dirs=[os.path.expandvars('$CUDA_HOME') + '/lib/x64', strBasepath],
+        libraries=['cudart', 'TH', 'THC', 'ATen']
 )

 if __name__ == '__main__':

src/HadamardProduct_cuda.cpp の作成（修正）

(1) src/HadamardProduct_cuda.c を src/HadamardProduct_cuda.cpp にrenameする。

> git mv src/HadamardProduct_cuda.c src/HadamardProduct_cuda.cpp

(2) src/HadamardProduct_cuda.cppを修正

--- a/src/HadamardProduct_cuda.cpp
+++ b/src/HadamardProduct_cuda.cpp
@@ -1,10 +1,13 @@
 #include <THC.h>
 #include <THCGeneral.h>
+#include <ATen/ATen.h>

 #include "HadamardProduct_kernel.h"

-extern THCState* state;
+//extern THCState* state;
+THCState *state = at::globalContext().thc_state;

+extern "C" {
 int HadamardProduct_cuda_forward(
        THCudaTensor* input1,
        THCudaTensor* input2,
@@ -37,4 +40,6 @@ int HadamardProduct_cuda_backward(
        );

        return 1;
+}
+
 }
\ No newline at end of file

install.batを作成

set CUDA_HOME=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v9.1

set TORCHCMD=python -c "import os; import torch; print(os.path.dirname(torch.__file__))"
set TORCH
for /f "usebackq delims=" %%a in (`%TORCHCMD%`) do set TORCH=%%a

gendef %TORCH%\lib\TH.dll
gendef %TORCH%\lib\THC.dll
gendef %TORCH%\lib\ATen.dll
lib /machine:x64 /def:TH.def 
lib /machine:x64 /def:THC.def 
lib /machine:x64 /def:ATen.def 

nvcc -c -o src/HadamardProduct_kernel.o src/HadamardProduct_kernel.cu --gpu-architecture=compute_52 --gpu-code=compute_52 --compiler-options -fPIC -I %TORCH%/lib/include/TH -I %TORCH%/lib/include/THC -Xcompiler "/MD /wd 4819"
python install.py

rename _ext\cunnex\_cunnex.so _cunnex.pyd

gendef.exeはmingw64に含まれるもの、lib.exeはVC2017に付属のものです。

ビルド

> install.bat

たくさんWarningがでますが、気にしないことにします。

テスト

> python test.py
0.0 <-- should be 0.0
True <-- should be true
0.0 <-- should be 0.0
True <-- should be true
0.0 <-- should be 0.0
True <-- should be true
switching to DataParallel mode
0.0 <-- should be 0.0
True <-- should be true
0.0 <-- should be 0.0
True <-- should be true
0.0 <-- should be 0.0
True <-- should be true

一応テストはパスするみたいです。