GPU/TPU で Mandelbrot (Google Colaboratory 編)

TensorFlow？

Google Colaboratory で TPU が使えるということで盛り上がっているようなので、訳の分からぬまま使ってみることにしてみました。Qiita 等各種記事を見る限り、TensorFlow というライブラリから呼び出すようになっているようです。

TensorFlow というものを調べてみると、どうも配列演算のサブルーチン集のようで、Fortran でいえば、’70年代や’80年代のアレイプロセッサ時代に逆戻りしたような趣きです。ただリダクション演算が拡張されていて、指定の任意の Rank 間でリダクション演算ができるようになっています。これは、Fortran で不満に思う所なので中々気が利くなと思いました。Tensor の名前の由来と思われます。

けしからんことに、各種言語用の API が用意されているのに、最もふさわしいと思われる Fortran 用がありません。死刑！ここでは、Python で利用してみます。

計算時間を比較すると、計算の性質かあるいは使い方を間違っているのか、GPU の方が TPU よりかなり速い感じでした。Mandelbrot 図形もなにか荒いので、計算間違ってるかもしれません。Python も初心者なので許してね。

Google 謹製例題 　

Hello, TPU in Colab

TensorFlow マニュアル

TensorFlow

プログラム

GPU

import tensorflow as tf
tf.test.gpu_device_name()

import numpy as np

N = 1024
Niter = 255
zero = tf.zeros([N, N])
one  = tf.fill([N, N], 1.0)
n = zero
x = np.linspace(-2.0, 1.0, N)
y = np.linspace(-1.5, 1.5, N)
xx = x[None, :]
yy = y[:, None]
c = tf.complex(xx, yy)
z = tf.multiply(c, 0.0)
for i in range(Niter):
    z = tf.add(tf.square(z), c)
    n = tf.add(n, tf.where(tf.abs(z) > 2.0, one, zero))

sess = tf.Session()

result = sess.run(n)

%matplotlib inline
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (20.0, 20.0)
plt.imshow(result)

実行例 GPU

TPU

import tensorflow as tf
import os

tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
print ('TPU address is', tpu_address)

with tf.Session(tpu_address) as session:
  devices = session.list_devices()

print ('TPU devices:')
devices

import numpy as np

N = 1024
Niter = 255
zero = tf.zeros([N, N])
one  = tf.fill([N, N], 1.0)
n = zero
x = np.linspace(-2.0, 1.0, N)
y = np.linspace(-1.5, 1.5, N)
xx = x[None, :]
yy = y[:, None]
c = tf.complex(xx, yy)
z = tf.multiply(c, 0.0)
for i in range(Niter):
    z = tf.add(tf.square(z), c)
    n = tf.add(n, tf.where(tf.abs(z) > 2.0, one, zero))

session = tf.Session(tpu_address)
try:
  print('Initializing...')
  session.run(tf.contrib.tpu.initialize_system())
  print('Running ops')
  result = session.run(n)
finally:
  # For now, TPU sessions must be shutdown separately from
  # closing the session.
  session.run(tf.contrib.tpu.shutdown_system())
  session.close()

%matplotlib inline
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (20.0, 20.0)
plt.imshow(result)

実行例 TPU

参考

Python での初期座標の作り方は、以下のページを参考にしました。
http://www.turbare.net/transl/scipy-lecture-notes/advanced/advanced_numpy/index.html

x = np.linspace(-2.0, 1.0, N)
y = np.linspace(-1.5, 1.5, N)
xx = x[None, :]
yy = y[:, None]