RADEON (ROCm) で PyTorch を使う。C++ API

Deep Learning フレームワークのほとんどが Python を使っています。C++ など他の言語も使えますが、Python と同じ使い方ができるものはあまり多くありません。

例えば CNTK の場合 Keras の backend として使うとこんな感じです。

# Keras
model= models.Sequential([
    layers.Dense( 64, input_shape=(784,), activation='relu' ),
    ~
  ])

CNTK の Python API を使った場合↓

# CNTK
xinput= cntk.input_variable( 784, np.float32 )
model= layers.Sequential([
    layers.Dense( 64, activation=cntk.relu ),
    ~
  ])

CNTK の C++ の場合 Parameter を確保して複数の Function に分解する必要があるのでこうなります。↓

// CNTK + C++
auto device= CNTK::DeviceDescriptor::UseDefaultDevice();
auto xinput= CNTK::InputVariable( { 784 }, CNTK::DataType::Float, L"xinput" );
auto weight= CNTK::Parameter( { 64, 784 }, CNTK::DataType::Float, CNTK::HeNormaliInitializer(), device );
auto bias= CNTK::Parameter( { 64 }, CNTK::DataType::Float, 0.0f, device );
auto x= CNTK::ReLU( CNTK::Plus( bias, CNTK::Times( weight, xinput ) ) );
~

PyTorch の場合は C++ でも Python と同じ書き方で同じ API なので覚えることが少なくて済みます。

// PyTorch C++
class ModelImpl : public torch::nn::Module {
    torch::nn::Linear  fc0;
public:
    ModelImpl()
    {
        fc0= register_module( "fc0", torch::nn::Linear( 784, 64 ) );
	~
    }
    torch::Tensor  forward( const torch::Tensor& x )
    {
        ~
        return  torch::relu( fc0( x ) );
    }
};
TORCH_MODULE(Model);

AMD RADEON を使用する場合の選択肢はいくつかあります。Python の場合は Keras + PlaidML が最も簡単で Windows/macOS でも使用できます。Linux の場合は ROCm が使えるため PlaidML 以外の選択肢が増えます。今回は C++ API を使いたいので Linux 上で ROCm を使用しました。

以下 RADEON で PyTorch (C++ API) を使うための作業メモです。

● ROCm の install

使用した環境は RADEON RX Vega 64 (gfx900) + Ubuntu 18.04 LTS です。

ROCm Documentation: Installation Guide

上のページを見ると Supported OS に 18.04.3 (Kernel 5.0) と書かれているため最初に Kernel を更新します。下記ページを参考にさせていただきました。

virtualiment: Ubuntu18.04にカーネル 5.0 をインストールする手順

カーネルの更新

$ sudo apt install --install-recommends linux-generic-hwe-18.04 xserver-xorg-hwe-18.04 
$ sudo reboot

次に手順通りに ROCm を入れます。こちらは下記のページを参考にさせていただきました。

ニートが始めるUE4開発日誌: RX470で機械学習ことはじめ その3 ~ROCmとTensorFlowのインストール~

install 前の更新とライブラリインストール

$ sudo apt update
$ sudo apt dist-upgrade
$ sudo apt install libnuma-dev
$ sudo reboot

リポジトリ追加と ROCm のインストール

$ wget http://repo.radeon.com/rocm/apt/debian/rocm.gpg.key
$ sudo apt-key add rocm.gpg.key
$ echo 'deb [arch=amd64] http://repo.radeon.com/rocm/apt/debian/ xenial main' | \
   sudo tee /etc/apt/sources.list.d/rocm.list
$ sudo apt install rocm-dkms
$ sudo usermod -a -G video $LOGNAME
$ sudo reboot

● PyTorch のビルド

PyTorch を ROCm でビルドします。ビルド手順は下記ページを参考にさせていただきました。

Lernapparat: Building PyTorch on ROCm

事前の準備

$ sudo apt install git cmake
$ sudo apt install python3-pip
$ pip3 install setuptools numpy

ビルドに必要な ROCm のパッケージを入れます。

$ sudo apt install rocm-libs miopen-hip rccl roctracer-dev

cmake のパッチのため下記のスクリプトを使います。

#!/bin/bash
for fn in $(find /opt/rocm/ -name \*.cmake ); do
  sudo sed --in-place='~' 's/find_dependency(hip)/find_dependency(HIP)/' $fn
done

上記の内容をファイル(replace_script.sh)に保存してから下記のように実行します。

$ bash ./replace_script.sh

PyTorch を checkout します。仮に展開場所を ~/pytorch とします。

$ git clone --recursive http://github.com/pytorch/pytorch

pytorch フォルダで下記のコマンドを実行します。

$ cd ~/pytorch
$ python3 tools/amd_build/build_amd.py

ビルドします。

$ cd ~/pytorch
$ RCCL_DIR=/opt/rocm/rccl/lib/cmake/rccl/ PYTORCH_ROCM_ARCH=gfx900 hip_DIR=/opt/rocm/hip/cmake/ USE_NVCC=OFF BUILD_CAFFE2_OPS=0 PATH=/usr/lib/ccache/:$PATH USE_CUDA=OFF python3 setup.py bdist_wheel

しばらくしたら Python 向けのパッケージが ~/pytorch/dist/*.whl にできます。C++ 向けの libtorch 相当が ~/pytorch/torch 以下になります。

include path は下記の通り。

~/pytorch/torch/include
~/pytorch/torch/include/torch/csrc/api/include

library は下記の場所に入るので、ビルド時の libpath と実行時の LD_LIBRARY_PATH に追加します。

~/pytorch/torch/lib

Link 時のライブラリ指定はこんな感じ。

-lc10 -ltorch -lc10_hip -ltorch_hip -ltorch_cpu -lc10

これで CUDA 向けのコードがそのまま RADEON で動くようになりました。ROCm (hip) でも PyTorch の API は cuda のままです。C++ API でも to( torch::kCUDA ) で OK。

import torch
for di in range(torch.cuda.device_count()):
    print( di, torch.cuda.get_device_name(di), torch.cuda.get_device_capability(di) )

例えば Python で↑のコードを実行すると次のように表示されます。

0 Vega 10 XT [Radeon RX Vega 64] (3, 0)

起動時に少々待たされますが起動してしまえば十分速いです。BatchSize は 64 以上がおすすめです。

関連ページ
HYPERでんち: Deep Learning

関連エントリ
Jetson Nano で TensorFlow の C 言語 API を使う