「PyTorch + XLA」のソースコード解析 (その1)

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

転職して、半年が経ちました。。。いやー、いろいろ仕事があって、オモロイですよ。

さて、本題です。

今日のブログは、Qiitaの「TensorFlow Advent Calendar 2018」の初日(12/1)の内容で、タイトルは、

「PyTorch + XLA」のソースコード解析

です。

「PyTorch + XLA」のソースコード解析は、
githubにて公開された翌日の2018年11月15日から11月18日の4日間で行いました。

今まで、TensorFlow XLA や TensorFlow Lite については、下記のようにGoogle Slide を利用して、PDFにして、Slideshareにて公開しましたが、今回は、Google Document にて、ソースコードを引用しながら、ソースコードを追っていくという感じにしました。

TensorFlow XLAは、中で何をやっているのか？、TensorFlow User Group ハード部 #2 (2017年4月21日)
TensorFlow XLAの可能性、Deep Learning Acceleration 勉強会(2017.09.03)@渋谷ヒカリエ
TensorFlow XLA とハードウェア、2017年9月30日(土)、Chainer Meetup #6
「ディープラーニングでは、エコシステムが大切よ！」、2018年2月16日(金)に開催された雑誌インターフェース2018年2月号オフ会
TensorFlow Lite (r1.5) & Android 8.1 Neural Network API、2018年3月7日(水)@LeapMind
Tensorflow dynamically loadable XLA plugin ソースコード解析、2018年5月24日(木)＠LeapMind
TensorFlow local Python XLA client、2018年8月4日Slideshareにて公開
TensorFlow XLA RPC、2018年8月4日Slideshareにて公開
Bridge TensorFlow to run on Intel nGraph backends (v0.4)、2018年8月25日Slideshareにて公開
Bridge TensorFlow to run on Intel nGraph backends (v0.5)、2018年8月25日Slideshareにて公開

このブログでは、作成した「Pytorch + XLA」のソースコード解析の

　0)、はじめに、
　1)、テストコード を見てみよう！
　8)、おわりに

をアップします。

残りの

　3)、PyTorchからXLAのOpへの変換
　4)、一旦、振り返り
　5)、実行環境の決定
　6)、ComputationClientでの実行
　7)、backward メソッド

については、別途、どのような形で公開するから、悩み中です。
興味がある方は、TwitterにMention or DM をくださいね。。(全員に返信するとは限りませんが、悪しからず)

TensorFlow User Group #9に申し込んでいますが、当選したら、そのときにでも。

==========================================================================================

PyTorch + XLAのソースコード解析

@Vengineer、2018.11.15 ～18

0)、はじめに

日本時間の2018年11月14日(埼玉県民の日)に、PyTorchからGoogle TPUを利用するためのソースコードが github にて公開されました。

https://github.com/pytorch/xla

PyTorchがGoogle Cloud TPUで動くということは、あたしの2018年10月9日のブログ、「PyTorchがTPUで動く」に書きました。このブログの中で紹介した「Google Cloud Blog「Introducing PyTorch across Google Cloud」の最後の方に、次のよう書いてありました。
引用します。

This prototype has successfully enabled us to train a PyTorch implementation of ResNet-50 on a Cloud TPU, and we’re planning to open source the prototype and then expand it in collaboration with the PyTorch community. Please email us at pytorch-tpu@googlegroups.com to tell us what types of PyTorch workloads you would be most interested in accelerating with Cloud TPUs!

ResNet-50 を Cloud TPU 上で動かすためのプロトタイプとありますね。それが今回公開されたソースコードかどうかはわかりませんが。。実際にソースコードを解析してみて、Resnet-50 を訓練(学習)できるようになっている感じです。

PyTorch の Model Zoo の Resnet-50 のコードを見てみたら、

　・nn.Conv2d
　・nn.BatchNorm2d
　・nn.ReLU
　・nn.MaxPool2d
　・nn.AdaptiveAvgPool2d
　・nn.Linear
　・x.view

があり、テストコードでは、

　・nn.Conv2d
　・nn.BatchNorm2d
　・F.relu
　・F.max_pool2d
　・F.avg_pool2d
　・nn.Linear
　・x.view

をサポートしています。なので、テストコードでテスト済みのレイヤーを使ったResnet-50は動くのでしょうね。

今回の実装では、推論 (foward) だけでなく、学習・訓練(forward & backward) もサポートしているようです。backward の方は、PyTorch の autodiff を上手に使っているようです。

さて、本命のGoogle TPUへのアクセスは、TensorFlow XLAにて行えますが、今回公開された「PyTorch + XLA」は、TensorFlow XLAを直接使うのではなく、XRT というものから XLA を呼び出しています。

README.md によると、次の3つのケースで利用できるようです。

ケース1)、local client を使って CPU を利用する場合 (XLAへの変換のデバッグ用かな？)

$ export XLA_USE_XRT=0 \
XLA_GRPC_HOST="" \
XLA_PLATFORM="CPU"

ケース2)、XRT client を使って、CPU を利用する場合 (XRTの練習みたいなもんかな？)

$ export XLA_USE_XRT=1 \
XRT_DEVICE_MAP="CPU:0;/job:localhost/replica:0/task:0/device:XLA_CPU:0" \
XRT_WORKERS="localhost:0;"

ケース3)、XRT client を使って、TPU を利用する場合 (本命！)

$ export XLA_USE_XRT=1 \
XRT_DEVICE_MAP="TPU:0;/job:tpu_worker/replica:0/task:0/device:TPU:0" \
XRT_WORKERS="tpu_worker:0;grpc://localhost:51000"

ちなみに、TensorFlow XRTについては、2018年9月4日の日記にアップしています。Cloud TPU用と書いてますね。正式には、TensorFlow r1.11 にてさらっと、リリースされています。

XRT の利用事例としては、Julia Computing, Inc.の Keno Fischerさんと Elliot Sabaさんの論文「Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs」がありますが、論文に多少のコードは載っているものの、残念ながらソースコードは公開されていません。

今回公開された「PyTorch + XLA」は、TensorFlow XRTを利用した唯一の公開ソースコードだと思います。
XLAではなく、XRTで利用できると、Cloud TPUを大量に利用できるのかな？と思っています。

以降、次のような内容について、ソースコードの解析をしていきます。

1)、テストコードを見てみよう！
2)、実装コードを見てみて
3)、PyTorchのレイヤーからTensorFlow XLAのOpへの変換
4)、一旦、振り返り
5)、実行環境の決定
6)、ComputationClientでの実行

1)、テストコードを見てみよう！

ソースコードを追う前に、まずは、テストコード ( test_operations.py ) を見てみましょう。
簡単な MulAdd のテストコードを以下に示します。

import torch_xla

class TestMulAdd(XlaTestCase):

def test(self):

class XlaMulAdd(nn.Module):

def forward(self, x, y):

return x * y + y

x = torch.rand(3, 5)
y = torch.rand(3, 5)
model = XlaMulAdd()
traced_model = torch.jit.trace(model, (x, y))
xla_model = torch_xla._C.XlaModule(traced_model)
inputs_xla = [torch_xla._C.XLATensor(x), torch_xla._C.XLATensor(y)]
output_xla = xla_model*1:

traced_model = torch.jit.trace(model, *input[0])
xla_model = torch_xla._C.XlaModule(

traced_model, use_full_conv_precision=True)

input_xla =
for n, replica_input in enumerate(input):

xla_replica_input =
for i in replica_input:

xla_replica_input.append(

torch_xla._C.XLATensor(i, '{}:{}'.format(device, n)))

input_xla.append(tuple(xla_replica_input))

output_xla = xla_model(*input_xla)
output =
for xla_replica_outputs in output_xla:

replica_outputs =
for o in _as_list(xla_replica_outputs):

replica_outputs.append(o.to_tensor())

output.append(tuple(replica_outputs))

return tuple(output)

else:

traced_model = torch.jit.trace(model, input)
xla_model = torch_xla._C.XlaModule(

traced_model, use_full_conv_precision=True)

input_xla = torch_xla._C.XLATensor(input)
output_xla = xla_model(tuple([input_xla]))
return output_xla[0][0].to_tensor()

else の下が TestMullAdd と同じようなコードになっています。

次は、もう少し凝った例 (と言っても、MNIST)。。

class XlaMNIST(nn.Module):

def __init__(self):

super(XlaMNIST, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)

def forward(self, x):

x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2(x), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x, dim=1)

class TestMNIST(XlaTestCase):

def test(self):

batch_size = 32
x = torch.randn(batch_size, 1, 28, 28)
model = XlaMNIST()
out = _xla_run(model, x)
expected = model(x)
self.assertEqualDbg(out.data, expected.data)

MulAdd とほぼ同じですね。これなら、それほど意識することなく、「PyTorch + XLA」を使って、Google Cloud TPU が利用できますね。

途中略。

8)、おわりに

　Facebook が PyTorch が Google の Cloud TPU で動く、と発表したときは、Googleは競争相手に自前のハードウェアアクセラレータである TPU を使うために協力したのか？と思いました。その後、よーく考えてみたら、GoogleはCloudベンダーであり、Facebookはユーザーであり、競争相手というよりはお客様だったんですよね。そう考えると、GoogleのCloud TPU を利用できる環境が増えるということはとってもいいことであるんですよね。

　何しろ、下記のように、GithubのStarsとForksの数では、PyTorchは、TensorFlow、Keras、Caffeに続いて、4番目に人気があるディープラーニングのフレームワークなんですから。
ちなみに、右側のStars増分とForks増分は、2017/11/04との比較によるものです。PyTorchは、この1年で100％以上の伸び率になっています。その次が、Keras、TensorFlowの順です。

今回の「PyTorch + XLA」のソースコード解析を行ったことにより、TensorFlow 以外のディープラーニング・フレームワークでも Google の Cloud TPU が利用できることが分かりました。

0)、はじめに、でも書きましたが、
Google Cloud Blog「Introducing PyTorch across Google Cloud」の最後の方に、次のよう書いてありました。

しかしながら、そのディープラーニング・フレームワークがCPUやGPU(大体がCUDA)で行っているすべてをCloud TPU上で実行できるようにするにはかなりの工数をかけないとできないと思います。それでも Facebook が Cloud TPU 上で PyTorch の Resnet-50のモデルを学習できるようにしたのにはそれなりの理由があったのだと思います。

最後に、ソースコードを公開してくれた Facebook および Google の皆さん、ありがとうございました。

＠Vengineer、2018.11.18

==========================================================================================

*1:tuple(inputs_xla)))
expected = model(x, y)
self.assertEqualDbg(output_xla[0][0].to_tensor().data, expected.data)

torch_xla を import しています。torch_xla._C の中の XlaModule, XLATensor を使っています。
XlaModule は、PyTorch のモデルを XLA 用に変換するようです。XLATensor は PythonのTensorをXLA側のTensorに変換しています。xla_model の入力引数に、XLATensor の inputs_xla を指定しています。計算結果の output_xla からは、output_xla[0][0].to_tensor().data にて、Pythonのデータに変換しています。

「PyTorch + XLA」 のソースコード解析

「PyTorch + XLA」のソースコード解析