AMDのGPUでディープラーニングのフレームワークの動作を検証

AMDのGPUでディープラーニングのワークフレームを動かす

近年、ディープラーニングが社会に浸透し、多くのシステムにおいて活用されています。そのディープラーニング技術を下支えしているのがGPUによる高速演算です。

ディープラーニングの演算を「学習」と「推論」の2つに分けたとき、特に学習において多くの計算資源が必要とされています。CPUで学習をしていてはとても待っていられないような計算も、GPUを活用することにより現実的な時間で終わらせることが可能になっています。

現在、ディープラーニングの学習をする際は、NVIDIA社のGPUを使用するのが主流となっています。データサイエンティスト・機械学習エンジニアの大半が、「GPUといえばNVIDIA」という認識でしょう。

しかし、NVIDIA社のGPUでしかディープラーニングはできないのでしょうか?本記事では、AMD社のGPUでディープラーニングのフレームワークを動かす方法と、実際に動いたことが確認できているソースコードを紹介します。

AMDのGPUでディープラーニングのフレームワークを動かす方法

まずは、どのようなソフトウェアを使うとAMDのGPUの上でディープラーニングのフレームワークが動かせるのかについて紹介します。

この記事は、実際に動いたソースコードの紹介を主眼に置いています。そのため、インストール手順を詳細に紹介するわけではありませんが、大まかな雰囲気だけ掴んで頂ければ幸いです。

AMDのGPUでディープラーニングのフレームワークは動くのか?

結論、動かすことは可能です。

もちろん、NVIDIAのGPUで動いているものがすべて動くと主張することはできませんが、多くのフレームワークがAMDのGPUでの動作をサポートする取り組みをしています。それを実現しているのが、ROCmというGPU演算用プラットフォームです。

厳密には少し異なりますが、イメージとしては、NVIDIAのGPUに対するCUDAが、AMDのGPUに対するROCmであると考えて頂くとわかりやすいでしょう。

2021年4月末現在の大きなニュースは、PyTorch 1.8から公式でROCmがサポートされるようになったことかと思います。PyTorchの公式サイトをご覧頂くと、CUDAやCPUと並列にROCm対応のPyTorchが提供されていることが確認できます。

大まかな流れとしては、

  1. ROCmのインストール
  2. ROCm対応版のディープラーニングフレームワークのインストール

という2つの操作で環境を適切に構築することにより、AMDのGPUでディープラーニングのフレームワークを動かすことが可能になります。

実行環境について

ROCm

2021年4月末現在、v4.1が最新バージョンです。しかし、PyTorchの対応状況を鑑みて、本ブログではv4.0での実行について紹介します。

ROCmのインストールは大変そうという印象をお持ちの方も多いと思いますが、実際にはコマンド数行で完了するように整備されています。

詳細は公式のインストールガイドをご覧ください。

対応GPU

ROCm4.0では、Vega 10、Vega 7nmおよびMI100アーキテクチャをサポートしています。今回の記事では、Vega 10アーキテクチャを採用するGPUであるRadeon RX Vega 56を使用します。

詳細は、公式のreadmeをご覧ください。

対応フレームワークについて

PyTorch、TensorFlow v1、TensorFlow v2の3フレームワークでは、手元での動作を確認できています。MXNet、JAXなど他のフレームワークでも対応が進められているようなので、次回以降の記事で紹介できたらと考えています。

PyTorch

先述の通り、バージョン1.8から公式にサポートされています。公式サイトの通りにコマンド入力するとインストール完了です。

TensorFlow

tensorflow-rocmという名前で、PyPIに登録されています。`pip install tensorflow-rocm==x.y.z`とすることにより、容易にインストールが可能です。

少し注意が必要なのは、メンテナンスバージョン(x.y.zと表記した時のz)でTensorFlowのバージョンとROCmのバージョンの対応を取っていることです。

例えば、tensorflow-rocm 2.4.0はROCm4.0向けのTensorFlow 2.4を、tensorflow-rocm 2.4.1はROCm4.1向けのTensorFlow 2.4を指しています。

詳細は、バージョン対応テーブルをご覧ください。

実行できることが確認できているサンプルコード

いくらカタログスペック的に動作するという主張を見ても、実際に動くかどうかはまた別の話という認識を多くのエンジニアが持っていると思います。筆者もそのうちの1人でした。

そこで、各フレームワークが提供しているサンプルコードが動いた実例を以下に紹介します。試してみたのは、次の3タスクです。

  • Image Classification(画像分類)
  • Object Detection(物体検出)
  • Semantic Segmentation(セマンティックセグメンテーション)

これらはすべて画像認識の基本的なタスクです。

以下のソースコードはまったく修正することなく、git cloneしてきた後実行コマンドを叩くのみで、正常に実行が完了することを確認できています。

PyTorch 1.8

Image Classification

オーソドックスなMNISTの分類タスクです。ネットワークは簡単なConvolutionとDropoutの組み合わせで、データセットはtorchvisionのものを使用しています。

Object Detection

Detectron2というObject Detection用フレームワークのサンプルコードが動くことも確認できています。ネットワークはMask-RCNNです。demo(推論)とtools(学習)の両方で動作を確認できています。

TensorFlow 1.15

Image Classification

データセットはDogs vs Cats、モデルはMobileNetV2のClassificationタスクです。ImageNetでpretrainしたモデルを初期値として学習するチュートリアルです。

TensorFlow  2.3

Image Classification

データセットはDogs vs Cats、モデルは簡単なConvolutionとDropoutの組み合わせのClassificationタスクです。

Semantic Segmentation

データセットはOxford-IIIT Pet Dataset、モデルはMovileNetV2をバックボーンにしたU-Netです。Semantic Segmentationタスクを実装しています。

実行できることが確認できているKaggleのSolution

サンプルコード以外に、KaggleのSolutionの実行も確認しました。

Kaggleとは、機械学習モデルの精度を競う大会(コンペ)のプラットフォームです。企業などから提供された問題に対し、世界中から集った参加者が予測精度を競います。

Solutionとは、そのコンペの解法、または解法を実装したソースコードだと理解して頂けると良いかと思います。

フレームワークのサンプルコードと比べて現実問題に近く、多くのリソースも必要とすることから検証対象として適切であると判断しました。

23rd place solution in Kaggle iMet Collection 2019

iMet Collection 2019というKaggleのコンペのSolutionです。

メトロポリタン美術館に所蔵されている美術品の画像に対し、適切なラベルを付けるという課題です。ラベルというのは、どのような美術品であるか、あるいはどのような文化に属しているかといったその美術品の属性情報です。

属性は一つの美術品に対し複数あることから(例えば、ある美術品は「イギリス」の「ライオン」の「銅像」で3ラベルとカウントします)、Multi ClassのImage Classificationとして解く問題でした。

フレームワークはPyTorchを採用しており、モデルはpretrainedmodelsのSE-ResNeXtなどを使用しています。このSolutionも、問題なく動作することを確認できました。

このSolutionの作者は筆者なのですが、自分がNVIDIAのGPUで開発したコードを修正することなくAMDのGPUで実行できるというのは意外でした。サンプルコードではなく、ある程度複雑なモデルの学習をAMDのGPUで実行できるというのはかなり活用の幅が広がるという印象です。

クラウド環境でAMDのGPUを使う方法

ここまで、動作が確認できているスクリプトを紹介してきました。上記のようなさまざまなスクリプトが問題なく実行できることから、ユースケースとコストによってはAMDのGPUクラウドを使用するという選択肢も生じてくるのではないでしょうか。

その選択肢を検討する際は、当社Morgenrotが提供する「M:CPP」のGPUクラウドがおすすめです。

MCPP

M:CPPでは、AMDのGPUを機械学習用にクラウド提供しています。安価な価格と、カスタマイズ性に優れていることが特徴です。研修・研究・PoC・本番システムなど、幅広い用途にご利用いただけます。

よりリーズナブルで、ハイエンドな環境を求めている場合には、お気軽にご相談ください。

まとめ

AMDのGPUに対するディープラーニングフレームワークのサポートはどんどん成熟してきています。本記事で紹介したように、Image Classification, Semantic Segmentation, Object Detectionといったメジャーなタスクの実行が確認できています。

ユースケースとコストに応じて、学習をAMDのGPUで実行することを検討に入れてみるのはいかがでしょうか?