TadaoYamaokaの開発日記

個人開発しているスマホアプリや将棋AIの開発ネタを中心に書いていきます。

将棋の合法手の数の統計

AlphaZero方式の強化学習では、指し手の確率分布を教師データとするため、局面の合法手の数の分だけ確率の値を保存しておく必要がある。
将棋の合法手の最大数は593であることが証明されているが、実際の対局ではそのような局面は現れない。
教師データを保存する容量を抑えるために、できれば実際に現れる局面の最大数に制限したい。
可変長フォーマットにするという案もあるが、機械学習でシャッフルやランダムサンプリングするため、固定長フォーマットにしておきたい。

そこで、実際の棋譜で合法手の数の統計量を調べてみた。

調査方法

以下のようなスクリプトでmatplotlibでヒストグラム表示と、Pandasで統計量を出力した。CSAファイルの解析にはcshogiを利用した。

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from cshogi import *
import sys
import glob

legal_moves = []

def read_kifu(kifu_list):
    positions = []
    parser = Parser()
    for filepath in kifu_list:
        parser.parse_csa_file(filepath.encode('utf-8'))
        board = Board()
        for move in parser.moves:
            legal_moves.append(len(board.leagal_move_list()))
            board.push(move)

kifu_list = glob.glob(sys.argv[1] + r"\*")

read_kifu(kifu_list)

plt.hist(legal_moves)
plt.show()

print(pd.Series(legal_moves).describe())
print(np.percentile(legal_moves, 99))
print(np.percentile(legal_moves, 99.9))

floodgateの1年分の統計

floodgateの2018年の棋譜を調査した結果は以下の通り。
f:id:TadaoYamaoka:20190218222632p:plain

count    4.754067e+06
mean     8.223365e+01
std      5.962379e+01
min      1.000000e+00
25%      3.800000e+01
50%      7.400000e+01
75%      1.110000e+02
max      4.730000e+02
99パーセンタイル 266
99.9パーセンタイル 347

棋譜数は47,813、局面の総数は4,754,067局面で、平均は82.2手で、最大は473手だった。
将棋の合法手の平均は約80手と言われているので、近い値になっている。
99.9パーセンタイルは347手で、ほとんどの局面で合法手の数は347手に収まっている。

elmoの自己対局で生成した教師データ

elmoの自己対局で生成した1000万局分でも調べてみた。
f:id:TadaoYamaoka:20190218223745p:plain

count    1.000000e+07
mean     1.021326e+02
std      6.950770e+01
min      1.000000e+00
25%      4.600000e+01
50%      9.600000e+01
75%      1.480000e+02
max      4.860000e+02
99パーセンタイル 284
99.9パーセンタイル 350

平均は102.1手、最大は486手、99.9パーセンタイルは350手だった。
平均がfloodgateの棋譜に比べて高いのは初期局面集を使っているためと思われる。

dlshogiの強化学習で生成した局面

dlshogiの強化学習で生成した250万局面でも調べてみた。
f:id:TadaoYamaoka:20190218224349p:plain

count    2.498567e+06
mean     9.635840e+01
std      6.277702e+01
min      2.000000e+00
25%      4.600000e+01
50%      8.900000e+01
75%      1.360000e+02
max      4.980000e+02
99パーセンタイル 275
99.9パーセンタイル 352

こちらも初期局面集を使用している。
平均は96.3手、最大は498手、99.9パーセンタイルは352手だった。

まとめ

調査した範囲で最大の合法手の数は、498手だった。
99.9パーセンタイルは352手のため、352に制限して、それを超えた分は、0に近い確率になっているはずなので除外しても問題なさそうである。
99パーセンタイルの284手に制限しても実用上問題ないかもしれない。

LeelaChessZeroではどうしているかソースを調べてみたところ、policyの出力のサイズでそのまま保存するという贅沢な使い方をしていた。
ディスクサイズとメモリサイズをケチるという発想はないようである。

Google ColabでAlphaZero Shogiのモデルを教師あり学習する

Google ColabでAlphaZero Shogiのモデルを論文に通り定義して、テストのために教師ありで学習してみました。
TPUでも学習して学習時間の比較もしてみました。

教師データには、elmoで生成したhcpe形式のデータを使用し、入力特徴量と正解ラベルの加工には、先日作成したPythonの将棋ライブラリ(cshogi)を使用しました。

モデルの定義

AlphaZeroの論文の通り、ResNetで、policyとvalueの2つの出力を持つネットワークを定義します。
ブロック数とフィルタ数、全結合層のユニット数はパラメータにしています。
policyの出力の畳み込みの後の位置ごとのバイアスはカスタムレイヤーを定義しています。

import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer, Input, Dense, Conv2D, BatchNormalization, Activation, Flatten, Add

class Bias(Layer):
    def __init__(self, **kwargs):
        super(Bias, self).__init__(**kwargs)

    def build(self, input_shape):
        self.b = self.add_weight(name='b',
                                 shape=(input_shape[1:]),
                                 initializer='zeros',
                                 trainable=True)
        super(Bias, self).build(input_shape)

    def call(self, x):
        return x + self.b

def conv_layer(inputs,
               filters,
               activation='relu',
               use_bias=True):

    x =  Conv2D(filters,
                kernel_size=3,
                strides=1,
                padding='same',
                data_format='channels_first',
                kernel_initializer='he_normal',
                use_bias=use_bias)(inputs)
    x = BatchNormalization(axis=1)(x)
    if activation is not None:
        x = Activation(activation)(x)
    return x

def ResNet(input_planes=45,
           res_blocks=20,
           filters=256,
           fcl_units=256,
           policy_planes=139):

    inputs = Input(shape=(input_planes, 9, 9))
    x = conv_layer(inputs, filters=filters, use_bias=False)

    for res_block in range(res_blocks):
        # bottleneck residual unit
        y = conv_layer(x, filters=filters, use_bias=False)
        y = conv_layer(y, filters=filters, use_bias=False, activation=None)
        x = Add()([x, y])
        x = Activation('relu')(x)

    # Add policy output
    policy_x = Conv2D(policy_planes,
                      kernel_size=1,
                      strides=1,
                      padding='same',
                      data_format='channels_first',
                      kernel_initializer='he_normal',
                      use_bias=False)(x)
    policy_x = Flatten()(policy_x)
    policy_y = Bias(name='policy')(policy_x)

    # Add value output
    value_x = conv_layer(x, filters=1)
    value_x = Flatten()(value_x)
    value_x = Dense(fcl_units,
                    activation='relu',
                    kernel_initializer='he_normal')(value_x)
    value_y = Dense(1,
                    activation='tanh',
                    kernel_initializer='he_normal',
                    name='value')(value_x)

    # Instantiate model
    return Model(inputs=inputs, outputs=[policy_y, value_y])

教師データ

教師データには、elmoを使用して、hcpe形式で生成したデータを使用します。
hcpe形式のデータには、

hcp ハフマン符号で圧縮した局面
bestMove16 指し手
gameResult 勝敗結果

が含まれます。

訓練データの局面数は、100万局面です。テストデータは10万局面です。

入力特徴量と正解データ

AlphaZeroの論文の通り、局面を入力特徴量に変換します。
ただし、hcpeには、局面の繰り返し数と手数が含まれないため、値を0にしています。
また、局面の繰り返し数は、単に繰り返しがあるかどうかだけ判定し、特徴を1面だけとしました。(今回は常に0)

持ち駒の数(prisoner count)は0から1の範囲の正規化を行わず、そのまま整数の値を設定しています。

policyの正解ラベルは、sparse_categorical_crossentropyを使用するので、one-hotベクトル化を行わず、整数で表しています。
valueの正解となる勝敗は、出力の活性関数にtanhを使うので、-1(負け),0(引き分け),1(勝ち)で表します。

履歴局面はなしです。

cshogiを使うと、以下のように記述できます。

board = Board()
def mini_batch(hcpes):
    features = np.zeros((len(hcpes), 45, 81), dtype=np.float32)
    action_labels = np.empty(len(hcpes), dtype=np.int)
    game_outcomes = np.empty(len(hcpes), dtype=np.float32)

    for i, hcpe in enumerate(hcpes):
        # Input features
        #   P1 piece 14
        #   P2 piece 14
        #   Repetitions 1
        #   P1 prisoner count 7
        #   P2 prisoner count 7
        #   Colour 1
        #   Total move count 1
        board.set_hcp(hcpe['hcp'])
        # piece
        pieces = board.pieces
        for sq in SQUARES:
            piece = pieces[sq]
            if piece != NONE:
                if piece >= WPAWN:
                    piece = piece - 2
                features[i][piece - 1][sq] = 1
        # repetition
        if board.is_draw() == REPETITION_DRAW:
            features[i][28].fill(1)
        # prisoner count
        pieces_in_hand = board.pieces_in_hand
        for c, hands in enumerate(pieces_in_hand):
            for hp, num in enumerate(hands):
                features[i][29 + c * 7 + hp].fill(num)
        # Colour
        if board.turn == WHITE:
            features[i][43].fill(1)
        # Total move count
        # not implement for learning from hcpe

        # Action representation
        #   Queen moves 64
        #   Knight moves 2
        #   Promoting queen moves 64
        #   Promoting knight moves 2
        #   Drop 7
        move = hcpe['bestMove16']
        if not move_is_drop(move):
            from_sq = move_from(move)
            to_sq = move_to(move)
            from_file, from_rank = divmod(from_sq, 9)
            to_file, to_rank = divmod(to_sq, 9)
            diff_file = to_file - from_file
            diff_rank = to_rank - from_rank
            if abs(diff_file) != 1 or abs(diff_rank) != 2:
                # Queen moves
                if diff_file < 0:
                    if diff_rank < 0:
                        move_dd = -diff_file - 1
                    elif diff_rank > 0:
                        move_dd = 8 - diff_file - 1
                    else:
                        move_dd = 16 - diff_file - 1
                elif diff_file > 0:
                    if diff_rank < 0:
                        move_dd = 24 + diff_file - 1
                    elif diff_rank > 0:
                        move_dd = 32 + diff_file - 1
                    else:
                        move_dd = 40 + diff_file - 1
                else:
                    if diff_rank < 0:
                        move_dd = 48 - diff_rank - 1
                    else:
                        move_dd = 56 + diff_rank - 1
            else:
                # Knight moves
                if diff_file < 0:
                    move_dd = 64
                else:
                    move_dd = 65

            promotion = 1 if move_is_promotion(move) else 0

            action_labels[i] = (promotion * 66 + move_dd) * 81 + from_sq
        else:
            # drop
            to_sq = move_to(move)
            hp = move_drop_hand_piece(move)
            action_labels[i] = (132 + hp) * 81 + to_sq

        # game outcome
        #   z: −1 for a loss, 0 for a draw, and +1 for a win
        gameResult = hcpe['gameResult']
        if board.turn == BLACK:
            if gameResult == BLACK_WIN:
                game_outcomes[i] = 1
            if gameResult == WHITE_WIN:
                game_outcomes[i] = -1
            else:
                game_outcomes[i] = 0
        else:
            if gameResult == BLACK_WIN:
                game_outcomes[i] = -1
            if gameResult == WHITE_WIN:
                game_outcomes[i] = 1
            else:
                game_outcomes[i] = 0

    return (features.reshape((len(hcpes), 45, 9, 9)), { 'policy': action_labels, 'value': game_outcomes })

Googleドライブのマウント

学習データは、Googleドライブにアップロードして、Google Colabからマウントしてアクセスします。

from google.colab import drive
drive.mount('/content/drive')

Googleドライブのファイルは、「drive/My Drive」からアクセスできます。

学習

hcpeデータはnumpyのfromfileでデータ形式を指定して読み込みます。
hcpeのデータ形式HuffmanCodedPosAndEvalはcshogiで定義しています。
入力特徴量はデータ量が大きいため、fit_generatorでミニバッチごとに入力特徴量と正解データを作成しながら学習します。

モデルのサイズは、10ブロック、192フィルタとしています。

import tensorflow as tf
from tensorflow.keras.optimizers import SGD
import numpy as np
import os
from cshogi import *

train_hcpe_path = 'drive/My Drive/hcpe/elmo_teacher_depth8_uniq-001-01'
test_hcpe_path = 'drive/My Drive/hcpe/elmo_teacher_depth8_uniq-test-01'
batchsize = 256
epochs = 1
weight_decay = 1e-4
use_tpu = True

model = ResNet(res_blocks=10, filters=192)

train_hcpes = np.fromfile(train_hcpe_path, dtype=HuffmanCodedPosAndEval)
test_hcpes = np.fromfile(test_hcpe_path, dtype=HuffmanCodedPosAndEval)

def datagen(hcpes, batchsize):
    while True:
        np.random.shuffle(hcpes)
        for i in range(0, len(hcpes) - batchsize, batchsize):
            yield mini_batch(hcpes[i:i+batchsize])

def categorical_crossentropy(y_true, y_pred):
    return tf.keras.backend.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)

def categorical_accuracy(y_true, y_pred):
    return tf.keras.metrics.sparse_categorical_accuracy(y_true, tf.nn.softmax(y_pred))

def binary_accuracy(y_true, y_pred):
    return tf.keras.metrics.binary_accuracy(tf.keras.backend.round((y_true + 1) / 2), y_pred, threshold=0)

# add weight decay
for layer in model.layers:
    if isinstance(layer, tf.keras.layers.Conv2D) or isinstance(layer, tf.keras.layers.Dense):
        layer.add_loss(tf.keras.regularizers.l2(weight_decay)(layer.kernel))

model.compile(optimizer=tf.train.MomentumOptimizer(learning_rate=0.01, momentum=0.9),
              loss={'policy': categorical_crossentropy, 'value': 'mse'},
              metrics={'policy': categorical_accuracy, 'value': binary_accuracy})

# TPU
if use_tpu:
    model = tf.contrib.tpu.keras_to_tpu_model(
        model,
        strategy=tf.contrib.tpu.TPUDistributionStrategy(
            tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
        )
    )

model.fit_generator(datagen(train_hcpes, batchsize), len(train_hcpes) // batchsize,
                    epochs=epochs,
                    validation_data=datagen(test_hcpes, batchsize), validation_steps=len(test_hcpes) // batchsize)

学習結果

TPUでの学習結果
INFO:tensorflow:Querying Tensorflow master (grpc://10.26.203.146:8470) for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 2190986923747320621)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 12416730924266005929)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 13813975489927709994)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 8536274068396247395)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 3398334431816235909)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 590228271416283825)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 15798692355657786920)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 1954696542070558698)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 898663154261538184)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 4503259854332690989)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 10060441942931197921)
WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.
INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(32,), dtype=tf.int32, name='core_id_40'), TensorSpec(shape=(32, 9, 9, 45), dtype=tf.float32, name='input_4_10'), TensorSpec(shape=(32, 1), dtype=tf.float32, name='policy_target_120'), TensorSpec(shape=(32, 1), dtype=tf.float32, name='value_target_120')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for input_4
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 38.74543237686157 secs
INFO:tensorflow:Setting weights on TPU model.
3905/3906 [============================>.] - ETA: 0s - loss: 4.4481 - policy_loss: 1.7086 - value_loss: 1.7085 - policy_categorical_accuracy: 1.7085 - value_binary_accuracy: 1.7084INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(32,), dtype=tf.int32, name='core_id_50'), TensorSpec(shape=(32, 9, 9, 45), dtype=tf.float32, name='input_4_10'), TensorSpec(shape=(32, 1), dtype=tf.float32, name='policy_target_120'), TensorSpec(shape=(32, 1), dtype=tf.float32, name='value_target_120')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for input_4
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 17.765841245651245 secs
390/390 [==============================] - 43s 110ms/step - loss: 3.7269 - policy_loss: 1.4889 - value_loss: 1.4889 - policy_categorical_accuracy: 1.4878 - value_binary_accuracy: 1.4872
3906/3906 [==============================] - 644s 165ms/step - loss: 4.4478 - policy_loss: 1.7085 - value_loss: 1.7084 - policy_categorical_accuracy: 1.7084 - value_binary_accuracy: 1.7083 - val_loss: 3.7269 - val_policy_loss: 1.4889 - val_value_loss: 1.4889 - val_policy_categorical_accuracy: 1.4878 - val_value_binary_accuracy: 1.4872
<tensorflow.python.keras.callbacks.History at 0x7f841a0954a8>

TPUでは、lossは正しく表示されますが、accuracyが正しく表示されませんでした。
ネットワークの出力が2つの場合に発生するようです。
まだKerasのTPUサポートは正式版ではないので、バグの可能性があります。

Google ColabのGPU(Tesla K80)での学習結果
WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/math_ops.py:3066: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.cast instead.
390/390 [==============================] - 62s 159ms/step - loss: 4.4728 - policy_loss: 3.4411 - value_loss: 0.2205 - policy_categorical_accuracy: 0.2172 - value_binary_accuracy: 0.7430
3906/3906 [==============================] - 1803s 462ms/step - loss: 5.2603 - policy_loss: 4.1786 - value_loss: 0.2296 - policy_categorical_accuracy: 0.1737 - value_binary_accuracy: 0.7433 - val_loss: 4.4728 - val_policy_loss: 3.4411 - val_value_loss: 0.2205 - val_policy_categorical_accuracy: 0.2172 - val_value_binary_accuracy: 0.7430
<tensorflow.python.keras.callbacks.History at 0x7f61b7ee5780>
ローカルPCのGPU(1080 Ti)での学習結果
WARNING:tensorflow:From C:\Anaconda3\lib\site-packages\tensorflow\python\ops\resource_variable_ops.py:435: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.
Instructions for updating:
Colocations handled automatically by placer.
WARNING:tensorflow:From C:\Anaconda3\lib\site-packages\tensorflow\python\keras\utils\losses_utils.py:170: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.cast instead.
WARNING:tensorflow:From C:\Anaconda3\lib\site-packages\tensorflow\python\ops\math_ops.py:3066: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.cast instead.
2019-02-17 16:51:40.039612: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2
2019-02-17 16:51:40.216090: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1433] Found device 0 with properties:
name: GeForce GTX 1080 Ti major: 6 minor: 1 memoryClockRate(GHz): 1.6705
pciBusID: 0000:01:00.0
totalMemory: 11.00GiB freeMemory: 9.10GiB
2019-02-17 16:51:40.224327: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 0
2019-02-17 16:51:40.638740: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:
2019-02-17 16:51:40.642624: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0
2019-02-17 16:51:40.645263: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N
2019-02-17 16:51:40.648587: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 8780 MB memory) -> physical GPU (device: 0, name: GeForce GTX 1080 Ti, pci bus id: 0000:01:00.0, compute capability: 6.1)
2019-02-17 16:51:45.858150: I tensorflow/stream_executor/dso_loader.cc:152] successfully opened CUDA library cublas64_100.dll locally
390/390 [==============================] - 16s 42ms/step - loss: 4.4274 - policy_loss: 3.3871 - value_loss: 0.2285 - policy_categorical_accuracy: 0.2272 - value_binary_accuracy: 0.7431
3906/3906 [==============================] - 451s 116ms/step - loss: 5.2415 - policy_loss: 4.1575 - value_loss: 0.2314 - policy_categorical_accuracy: 0.1760 - value_binary_accuracy: 0.7424 - val_loss: 4.4274 - val_policy_loss: 3.3871 - val_value_loss: 0.2285 - val_policy_categorical_accuracy: 0.2272 - val_value_binary_accuracy: 0.7431

GPUでは、accuracyも正しく表示されています。

比較

TPUとGPUで学習時間を比較した結果は以下の通りでした。

TPU 644s
Google ColabのGPU(Tesla K80) 1803s
ローカルPCのGPU(1080Ti) 451s

バッチサイズを2048に増やすと以下の通りになりました。(Google ColabのGPUは遅いので測定から外しています。)

TPU 806s
ローカルPCのGPU(1080Ti) 516s

20ブロック、256フィルタ、バッチサイズ256では以下の通りでした。

TPU 1448s
ローカルPCのGPU(1080Ti) 1267s

いずれの条件でも、ローカルPCのGPU(1080Ti)が早いという結果になりました。
条件を変えると変わってくるかもしれません。

Google Colabを使う場合は、TPUの方がGPUより圧倒的に速いです。
ただし、TPUはまだKerasでの動作が安定していないため、安心して使えるのはTensorFlow 2.0で正式サポートされてからになりそうです。

NHWC vs NCHW on Google Colab

畳み込みの入力データの形式には、NHWCとNCHW があるが、どちらがTPUに最適か実験してみた。

TensorFlowのデフォルトはNHWCで、ChainerのデフォルトはNCHWになっている。

cuDNNはNCHWに最適化されている。
https://www.tensorflow.org/guide/performance/overview

しかし、TensorCoreは、NHWCに最適化されている。
Volta Tensor コア GPU が AI パフォーマンスの新記録を達成 | NVIDIA

なお、TensorFlowのCPUモードでは、MaxpoolingはNHWCしかサポートしていない。

TPUではどちらに最適化されているか調べてもわからなかった。

そこで、カラー画像のデータセットであるCIFAR-10を使って、2層の畳み込み層のあるネットワークを学習させて、実際に測ってみた。

NHWC on TPU

import tensorflow as tf
import os

(x_train, y_train),(x_test, y_test) = tf.keras.datasets.cifar10.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Conv2D(input_shape=(32, 32, 3), filters=256, kernel_size=3, padding='same', activation=tf.nn.relu),
  tf.keras.layers.MaxPool2D(),
  tf.keras.layers.Conv2D(filters=256, kernel_size=3, padding='same', activation=tf.nn.relu),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dense(10)
])

def loss(y_true, y_pred):
    return tf.keras.backend.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)

def accuracy(y_true, y_pred):
    return tf.keras.metrics.sparse_categorical_accuracy(y_true, tf.nn.softmax(y_pred))

# TPU
model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
    )
)

model.compile(optimizer='adam',
              loss=loss,
              metrics=[accuracy])
  
model.fit(x_train, y_train, batch_size=1024, epochs=5)
model.evaluate(x_test, y_test)
実行結果
INFO:tensorflow:Querying Tensorflow master (grpc://10.38.219.210:8470) for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 14236975761588034518)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 18323984429862649922)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 8554836709659081505)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 14275939449935473668)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 12561286578595138955)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 4788377722856588897)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 6433296346626590566)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 17630926159361266221)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 1107264524916151670)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 9476024643346962673)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 5663133602415764923)
WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.
Epoch 1/5
INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(128,), dtype=tf.int32, name='core_id_40'), TensorSpec(shape=(128, 32, 32, 3), dtype=tf.float32, name='conv2d_8_input_10'), TensorSpec(shape=(128, 1), dtype=tf.float32, name='dense_7_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for conv2d_8_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7f16729f7390> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 8.377423286437988 secs
INFO:tensorflow:Setting weights on TPU model.
INFO:tensorflow:CPU -> TPU lr: 0.0010000000474974513 {0.001}
INFO:tensorflow:CPU -> TPU beta_1: 0.8999999761581421 {0.9}
INFO:tensorflow:CPU -> TPU beta_2: 0.9990000128746033 {0.999}
INFO:tensorflow:CPU -> TPU decay: 0.0 {0.0}
WARNING:tensorflow:Cannot update non-variable config: epsilon
WARNING:tensorflow:Cannot update non-variable config: amsgrad
48128/50000 [===========================>..] - ETA: 1s - loss: 1.9748 - accuracy: 0.2997INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(106,), dtype=tf.int32, name='core_id_40'), TensorSpec(shape=(106, 32, 32, 3), dtype=tf.float32, name='conv2d_8_input_10'), TensorSpec(shape=(106, 1), dtype=tf.float32, name='dense_7_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for conv2d_8_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7f16729f7390> [<tf.Variable 'tpu_139734392716536/Adam/iterations:0' shape=() dtype=int64>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f167229c898>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f16722451d0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1672245780>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f16721fea58>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f16721c5c50>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f167218fbe0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f16720fde10>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f16720c4668>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f16720903c8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1672059b70>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1672025d68>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671f909e8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671f5bf98>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671f213c8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671eec198>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671eb1cf8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671dfcdd8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671dc3e10>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671db2b70>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671cfc7b8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671cc5d68>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671cb2f60>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671bfaa90>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f1671bc66d8>]
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 8.441371440887451 secs
50000/50000 [==============================] - 43s 850us/sample - loss: 1.9587 - accuracy: 0.3050
Epoch 2/5
50000/50000 [==============================] - 7s 137us/sample - loss: 1.5943 - accuracy: 0.4338
Epoch 3/5
50000/50000 [==============================] - 7s 138us/sample - loss: 1.8239 - accuracy: 0.4031
Epoch 4/5
50000/50000 [==============================] - 7s 135us/sample - loss: 1.4532 - accuracy: 0.4980
Epoch 5/5
50000/50000 [==============================] - 7s 136us/sample - loss: 1.2632 - accuracy: 0.5597
INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(4,), dtype=tf.int32, name='core_id_50'), TensorSpec(shape=(4, 32, 32, 3), dtype=tf.float32, name='conv2d_8_input_10'), TensorSpec(shape=(4, 1), dtype=tf.float32, name='dense_7_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for conv2d_8_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7f167066ac88> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 4.571932315826416 secs
 9952/10000 [============================>.] - ETA: 0s - loss: 2.0134 - accuracy: 0.3715INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(2,), dtype=tf.int32, name='core_id_50'), TensorSpec(shape=(2, 32, 32, 3), dtype=tf.float32, name='conv2d_8_input_10'), TensorSpec(shape=(2, 1), dtype=tf.float32, name='dense_7_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for conv2d_8_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7f167066ac88> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 3.0498554706573486 secs
10000/10000 [==============================] - 14s 1ms/sample - loss: 2.0130 - accuracy: 0.3718
[2.0130336515426634, 0.3718]

NCHW on TPU

import tensorflow as tf
import numpy as np
import os

(x_train, y_train),(x_test, y_test) = tf.keras.datasets.cifar10.load_data()
x_train, x_test = np.transpose(x_train, [0, 3, 1, 2]), np.transpose(x_test, [0, 3, 1, 2])
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Conv2D(input_shape=(3, 32, 32), filters=256, kernel_size=3, padding='same', activation=tf.nn.relu, data_format='channels_first'),
  tf.keras.layers.MaxPool2D(data_format='channels_first'),
  tf.keras.layers.Conv2D(filters=256, kernel_size=3, padding='same', activation=tf.nn.relu, data_format='channels_first'),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dense(10)
])

def loss(y_true, y_pred):
    return tf.keras.backend.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)

def accuracy(y_true, y_pred):
    return tf.keras.metrics.sparse_categorical_accuracy(y_true, tf.nn.softmax(y_pred))

# TPU
model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
    )
)

model.compile(optimizer='adam',
              loss=loss,
              metrics=[accuracy])
  
model.fit(x_train, y_train, batch_size=1024, epochs=5)
model.evaluate(x_test, y_test)
実行結果
INFO:tensorflow:Querying Tensorflow master (grpc://10.38.219.210:8470) for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 14236975761588034518)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 18323984429862649922)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 8554836709659081505)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 14275939449935473668)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 12561286578595138955)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 4788377722856588897)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 6433296346626590566)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 17630926159361266221)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 1107264524916151670)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 9476024643346962673)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 5663133602415764923)
WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.
Epoch 1/5
INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(128,), dtype=tf.int32, name='core_id_70'), TensorSpec(shape=(128, 3, 32, 32), dtype=tf.float32, name='conv2d_12_input_10'), TensorSpec(shape=(128, 1), dtype=tf.float32, name='dense_11_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for conv2d_12_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7f166d50c518> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 16.735506296157837 secs
INFO:tensorflow:Setting weights on TPU model.
INFO:tensorflow:CPU -> TPU lr: 0.0010000000474974513 {0.001}
INFO:tensorflow:CPU -> TPU beta_1: 0.8999999761581421 {0.9}
INFO:tensorflow:CPU -> TPU beta_2: 0.9990000128746033 {0.999}
INFO:tensorflow:CPU -> TPU decay: 0.0 {0.0}
WARNING:tensorflow:Cannot update non-variable config: epsilon
WARNING:tensorflow:Cannot update non-variable config: amsgrad
48128/50000 [===========================>..] - ETA: 1s - loss: 2.0741 - accuracy: 0.2895INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(106,), dtype=tf.int32, name='core_id_70'), TensorSpec(shape=(106, 3, 32, 32), dtype=tf.float32, name='conv2d_12_input_10'), TensorSpec(shape=(106, 1), dtype=tf.float32, name='dense_11_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for conv2d_12_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7f166d50c518> [<tf.Variable 'tpu_139734303609520/Adam/iterations:0' shape=() dtype=int64>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cdb1c50>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cd53518>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cd53ac8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cd125f8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166ccd9e80>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cc460f0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cc10278>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cbdc8d0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cba36a0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cb6b710>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166cab64e0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166caa3550>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166ca6c2e8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c9b3b38>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c97d828>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c947be0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c8b5e10>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c879b00>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c849748>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c80ef98>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c780ef0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c749a20>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c70f668>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7f166c6d8eb8>]
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 18.13627290725708 secs
50000/50000 [==============================] - 79s 2ms/sample - loss: 2.0549 - accuracy: 0.2950
Epoch 2/5
50000/50000 [==============================] - 6s 129us/sample - loss: 1.6440 - accuracy: 0.4250
Epoch 3/5
50000/50000 [==============================] - 7s 130us/sample - loss: 1.5291 - accuracy: 0.4633
Epoch 4/5
50000/50000 [==============================] - 6s 129us/sample - loss: 1.5411 - accuracy: 0.4643
Epoch 5/5
50000/50000 [==============================] - 6s 129us/sample - loss: 1.4702 - accuracy: 0.4918
INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(4,), dtype=tf.int32, name='core_id_80'), TensorSpec(shape=(4, 3, 32, 32), dtype=tf.float32, name='conv2d_12_input_10'), TensorSpec(shape=(4, 1), dtype=tf.float32, name='dense_11_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for conv2d_12_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7f166b0f7c88> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 11.619062662124634 secs
 9952/10000 [============================>.] - ETA: 0s - loss: 1.9147 - accuracy: 0.3833INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(2,), dtype=tf.int32, name='core_id_80'), TensorSpec(shape=(2, 3, 32, 32), dtype=tf.float32, name='conv2d_12_input_10'), TensorSpec(shape=(2, 1), dtype=tf.float32, name='dense_11_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for conv2d_12_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7f166b0f7c88> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 6.101793527603149 secs
10000/10000 [==============================] - 27s 3ms/sample - loss: 1.9154 - accuracy: 0.3828
[1.91538902759552, 0.38279998]

比較

1 2 3 4 5 eval.
NHWC 43s 7s 7s 7s 7s 14s
NCHW 79s 6s 7s 6s 6s 27s

NCHWは1エポック目が遅いが、2エポック目からわずかに速くなっている。
inferenceはNHWCが速い。これもNCHWの初回の遅さがネックになっていると思われる。

TPUはNCHWは初回の時間がかかるが、学習データが多い場合は、最終的には速くなりそうである。

追試

GPUでも比較してみた。

1 2 3 4 5 eval.
NHWC 24s 23s 23s 23s 23s 3s
NCHW 22s 22s 23s 22s 22s 3s

GPUでもNCHWがわずかに速い。

追試2

ローカルのPCの1080Tiでも測ってみた。

1 2 3 4 5 eval.
NHWC 10s 8s 8s 8s 8s 1s
NCHW 9s 8s 7s 7s 7s 1s

やはりNCHWがわずかに速い。

Google ColabでTPUを使ってみる

ほぼ自分用のメモです。

Google Colabで、Kerasを使ってTPUでMNISTの学習を試してみた。

TPUを有効にするには、「ランタイムのタイプを変更」からハードウェアアクセラレータを「TPU」に変更する必要がある。

KerasでTPUでMNISTを学習するには以下のように記述する。

import tensorflow as tf
import os

mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dense(10)
])

def sparse_categorical_crossentropy(y_true, y_pred):
    return tf.keras.backend.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)

def sparse_categorical_accuracy(y_true, y_pred):
    return tf.keras.metrics.sparse_categorical_accuracy(y_true, tf.nn.softmax(y_pred))

model.compile(optimizer='adam',
              loss=sparse_categorical_crossentropy,
              metrics=[sparse_categorical_accuracy])

# TPU
tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
strategy = tf.contrib.tpu.TPUDistributionStrategy(tpu_cluster_resolver)
model = tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)
    
model.fit(x_train, y_train, batch_size=1024, epochs=5)
model.evaluate(x_test, y_test)

TPUのコメントがある部分で、CPU/GPUのモデルからTPUのモデルに変換している。

実行結果

INFO:tensorflow:Querying Tensorflow master (grpc://10.58.99.162:8470) for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 12055453447884746007)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 18058977096921453459)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 14808883871349854203)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 12787960678777081298)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 12370191287575444036)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 10061541009258236286)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 18415570492467594980)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 6989406478957311628)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 5293577435862379493)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 11465669036580936775)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 13652277581393476684)
WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
Epoch 1/5
INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(128,), dtype=tf.int32, name='core_id_20'), TensorSpec(shape=(128, 28, 28), dtype=tf.float32, name='flatten_6_input_10'), TensorSpec(shape=(128, 1), dtype=tf.float32, name='dense_13_target_30')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for flatten_6_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fe1974d4c88> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 2.0240049362182617 secs
INFO:tensorflow:Setting weights on TPU model.
INFO:tensorflow:CPU -> TPU lr: 0.0010000000474974513 {0.001}
INFO:tensorflow:CPU -> TPU beta_1: 0.8999999761581421 {0.9}
INFO:tensorflow:CPU -> TPU beta_2: 0.9990000128746033 {0.999}
INFO:tensorflow:CPU -> TPU decay: 0.0 {0.0}
WARNING:tensorflow:Cannot update non-variable config: epsilon
WARNING:tensorflow:Cannot update non-variable config: amsgrad
57344/60000 [===========================>..] - ETA: 0s - loss: 0.5631 - sparse_categorical_accuracy: 0.8493INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(76,), dtype=tf.int32, name='core_id_20'), TensorSpec(shape=(76, 28, 28), dtype=tf.float32, name='flatten_6_input_10'), TensorSpec(shape=(76, 1), dtype=tf.float32, name='dense_13_target_30')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for flatten_6_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fe1974d4c88> [<tf.Variable 'tpu_140606887606032/Adam/iterations:0' shape=() dtype=int64>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe19708d278>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe19708dc18>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe19708de80>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe19706cfd0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe196fb5f60>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe196f7e278>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe196f49198>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe196f10cf8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe196ed89e8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe196ea4e10>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe196e11b70>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe196dda7b8>]
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 1.968165397644043 secs
60000/60000 [==============================] - 13s 209us/sample - loss: 0.5511 - sparse_categorical_accuracy: 0.8522
Epoch 2/5
60000/60000 [==============================] - 2s 29us/sample - loss: 0.2344 - sparse_categorical_accuracy: 0.9342
Epoch 3/5
60000/60000 [==============================] - 2s 32us/sample - loss: 0.1832 - sparse_categorical_accuracy: 0.9477
Epoch 4/5
60000/60000 [==============================] - 2s 31us/sample - loss: 0.1496 - sparse_categorical_accuracy: 0.9579
Epoch 5/5
60000/60000 [==============================] - 2s 27us/sample - loss: 0.1225 - sparse_categorical_accuracy: 0.9654
INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(4,), dtype=tf.int32, name='core_id_30'), TensorSpec(shape=(4, 28, 28), dtype=tf.float32, name='flatten_6_input_10'), TensorSpec(shape=(4, 1), dtype=tf.float32, name='dense_13_target_30')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for flatten_6_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fe196162898> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 1.5682730674743652 secs
 9888/10000 [============================>.] - ETA: 0s - loss: 0.1541 - sparse_categorical_accuracy: 0.9534INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(2,), dtype=tf.int32, name='core_id_30'), TensorSpec(shape=(2, 28, 28), dtype=tf.float32, name='flatten_6_input_10'), TensorSpec(shape=(2, 1), dtype=tf.float32, name='dense_13_target_30')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for flatten_6_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fe196162898> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 1.0298957824707031 secs
10000/10000 [==============================] - 7s 655us/sample - loss: 0.1553 - sparse_categorical_accuracy: 0.9530
[0.15529378306418656, 0.95299995]

同じモデルをGPUで学習すると、

Epoch 1/5
60000/60000 [==============================] - 1s 11us/sample - loss: 0.5526 - sparse_categorical_accuracy: 0.8556
Epoch 2/5
60000/60000 [==============================] - 1s 10us/sample - loss: 0.2269 - sparse_categorical_accuracy: 0.9368
Epoch 3/5
60000/60000 [==============================] - 1s 9us/sample - loss: 0.1711 - sparse_categorical_accuracy: 0.9521
Epoch 4/5
60000/60000 [==============================] - 1s 10us/sample - loss: 0.1357 - sparse_categorical_accuracy: 0.9625
Epoch 5/5
60000/60000 [==============================] - 1s 9us/sample - loss: 0.1116 - sparse_categorical_accuracy: 0.9693
10000/10000 [==============================] - 1s 70us/sample - loss: 0.1123 - sparse_categorical_accuracy: 0.9671
[0.11232251176461577, 0.9671]

TPUよりGPUの方が早いという結果だった。

畳み込みのモデルを学習

今度は、畳み込み層のあるモデルに変更して学習してみた。

import tensorflow as tf
import os

mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
x_train, x_test = x_train.reshape(x_train.shape[0], 28, 28, 1), x_test.reshape(x_test.shape[0], 28, 28, 1)

model = tf.keras.models.Sequential([
  tf.keras.layers.Conv2D(input_shape=(28, 28, 1), filters=256, kernel_size=3, padding='same', activation=tf.nn.relu),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dense(10)
])

def sparse_categorical_crossentropy(y_true, y_pred):
    return tf.keras.backend.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)

def sparse_categorical_accuracy(y_true, y_pred):
    return tf.keras.metrics.sparse_categorical_accuracy(y_true, tf.nn.softmax(y_pred))

model.compile(optimizer='adam',
              loss=sparse_categorical_crossentropy,
              metrics=[sparse_categorical_accuracy])

# TPU
tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
strategy = tf.contrib.tpu.TPUDistributionStrategy(tpu_cluster_resolver)
model = tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)
    
model.fit(x_train, y_train, batch_size=1024, epochs=5)
model.evaluate(x_test, y_test)
実行結果
INFO:tensorflow:Querying Tensorflow master (grpc://10.58.99.162:8470) for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 12055453447884746007)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 18058977096921453459)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 14808883871349854203)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 12787960678777081298)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 12370191287575444036)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 10061541009258236286)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 18415570492467594980)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 6989406478957311628)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 5293577435862379493)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 11465669036580936775)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 13652277581393476684)
WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
Epoch 1/5
INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(128,), dtype=tf.int32, name='core_id_40'), TensorSpec(shape=(128, 28, 28, 1), dtype=tf.float32, name='conv2d_input_10'), TensorSpec(shape=(128, 1), dtype=tf.float32, name='dense_15_target_30')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for conv2d_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fe1958053c8> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 11.533315420150757 secs
INFO:tensorflow:Setting weights on TPU model.
INFO:tensorflow:CPU -> TPU lr: 0.0010000000474974513 {0.001}
INFO:tensorflow:CPU -> TPU beta_1: 0.8999999761581421 {0.9}
INFO:tensorflow:CPU -> TPU beta_2: 0.9990000128746033 {0.999}
INFO:tensorflow:CPU -> TPU decay: 0.0 {0.0}
WARNING:tensorflow:Cannot update non-variable config: epsilon
WARNING:tensorflow:Cannot update non-variable config: amsgrad
58368/60000 [============================>.] - ETA: 0s - loss: 0.3569 - sparse_categorical_accuracy: 0.8933INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(76,), dtype=tf.int32, name='core_id_40'), TensorSpec(shape=(76, 28, 28, 1), dtype=tf.float32, name='conv2d_input_10'), TensorSpec(shape=(76, 1), dtype=tf.float32, name='dense_15_target_30')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for conv2d_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fe1958053c8> [<tf.Variable 'tpu_140606857389224/Adam/iterations:0' shape=() dtype=int64>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe195210828>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe1951b6160>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe1951b6710>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe1951f2940>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe195139a90>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe195103f28>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe1950cc4a8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe195039550>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe1950032e8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe194fcab38>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe194f91828>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe194f5fbe0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe194ecbe10>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe194e94b00>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe194e5e748>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe194e25f98>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe194d94f28>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fe194d5da90>]
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 11.685662269592285 secs
60000/60000 [==============================] - 43s 718us/sample - loss: 0.3501 - sparse_categorical_accuracy: 0.8954
Epoch 2/5
60000/60000 [==============================] - 4s 68us/sample - loss: 0.0625 - sparse_categorical_accuracy: 0.9818
Epoch 3/5
60000/60000 [==============================] - 4s 65us/sample - loss: 0.0338 - sparse_categorical_accuracy: 0.9898
Epoch 4/5
60000/60000 [==============================] - 4s 68us/sample - loss: 0.0195 - sparse_categorical_accuracy: 0.9947
Epoch 5/5
60000/60000 [==============================] - 4s 68us/sample - loss: 0.0134 - sparse_categorical_accuracy: 0.9962
INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(4,), dtype=tf.int32, name='core_id_50'), TensorSpec(shape=(4, 28, 28, 1), dtype=tf.float32, name='conv2d_input_10'), TensorSpec(shape=(4, 1), dtype=tf.float32, name='dense_15_target_30')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for conv2d_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fe193c579e8> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 3.3587257862091064 secs
 9824/10000 [============================>.] - ETA: 0s - loss: 0.0863 - sparse_categorical_accuracy: 0.9748INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(2,), dtype=tf.int32, name='core_id_50'), TensorSpec(shape=(2, 28, 28, 1), dtype=tf.float32, name='conv2d_input_10'), TensorSpec(shape=(2, 1), dtype=tf.float32, name='dense_15_target_30')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for conv2d_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fe193c579e8> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 2.2662925720214844 secs
10000/10000 [==============================] - 11s 1ms/sample - loss: 0.0855 - sparse_categorical_accuracy: 0.9749
[0.08545259298430756, 0.97489995]

同じモデルをGPUで学習すると、

Epoch 1/5
60000/60000 [==============================] - 33s 551us/sample - loss: 0.3539 - sparse_categorical_accuracy: 0.8934
Epoch 2/5
60000/60000 [==============================] - 32s 540us/sample - loss: 0.0571 - sparse_categorical_accuracy: 0.9834
Epoch 3/5
60000/60000 [==============================] - 32s 540us/sample - loss: 0.0286 - sparse_categorical_accuracy: 0.9922
Epoch 4/5
60000/60000 [==============================] - 32s 541us/sample - loss: 0.0152 - sparse_categorical_accuracy: 0.9962
Epoch 5/5
60000/60000 [==============================] - 33s 546us/sample - loss: 0.0087 - sparse_categorical_accuracy: 0.9980
10000/10000 [==============================] - 4s 396us/sample - loss: 0.0458 - sparse_categorical_accuracy: 0.9845
[0.04582856161564123, 0.9845]

今度は、TPUの方が圧倒的に速くなった。
モデルの計算量が多いほどTPUが速いようだ。

2019/2/13

公式のページの方法に書き直した。
https://www.tensorflow.org/guide/using_tpuwww.tensorflow.org

import tensorflow as tf
import os

mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
x_train, x_test = x_train.reshape(x_train.shape[0], 28, 28, 1), x_test.reshape(x_test.shape[0], 28, 28, 1)

model = tf.keras.models.Sequential([
  tf.keras.layers.Conv2D(input_shape=(28, 28, 1), filters=256, kernel_size=3, padding='same', activation=tf.nn.relu),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dense(10)
])

def sparse_categorical_crossentropy(y_true, y_pred):
    return tf.keras.backend.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)

def sparse_categorical_accuracy(y_true, y_pred):
    return tf.keras.metrics.sparse_categorical_accuracy(y_true, tf.nn.softmax(y_pred))

# TPU
model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
    )
)
   
model.compile(optimizer='adam',
              loss=sparse_categorical_crossentropy,
              metrics=[sparse_categorical_accuracy])

model.fit(x_train, y_train, batch_size=1024, epochs=5)
model.evaluate(x_test, y_test)
実行結果
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 2593207894788626939)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 13250067999102773720)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 4557324841645948994)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 17679983793927275190)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 2642697500206372799)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 12155722493614205727)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 16776203674625333233)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 219400119950085365)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 12315577173932374116)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 8568424978166959712)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 11493405250562256329)
WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.
Epoch 1/5
INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(128,), dtype=tf.int32, name='core_id_20'), TensorSpec(shape=(128, 28, 28, 1), dtype=tf.float32, name='conv2d_1_input_10'), TensorSpec(shape=(128, 1), dtype=tf.float32, name='dense_3_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for conv2d_1_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fec0d48f198> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 10.637016296386719 secs
INFO:tensorflow:Setting weights on TPU model.
INFO:tensorflow:CPU -> TPU lr: 0.0010000000474974513 {0.001}
INFO:tensorflow:CPU -> TPU beta_1: 0.8999999761581421 {0.9}
INFO:tensorflow:CPU -> TPU beta_2: 0.9990000128746033 {0.999}
INFO:tensorflow:CPU -> TPU decay: 0.0 {0.0}
WARNING:tensorflow:Cannot update non-variable config: epsilon
WARNING:tensorflow:Cannot update non-variable config: amsgrad
58368/60000 [============================>.] - ETA: 0s - loss: 0.3567 - sparse_categorical_accuracy: 0.8953INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(76,), dtype=tf.int32, name='core_id_20'), TensorSpec(shape=(76, 28, 28, 1), dtype=tf.float32, name='conv2d_1_input_10'), TensorSpec(shape=(76, 1), dtype=tf.float32, name='dense_3_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for conv2d_1_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fec0d48f198> [<tf.Variable 'tpu_140651824059952/Adam/iterations:0' shape=() dtype=int64>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cea81d0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cea8ac8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0ce43278>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0ce070b8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cdd2240>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cd9a898>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cd5f668>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cd2a6d8>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0ccf3f60>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cc639b0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cc28358>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cbf29b0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cbbb780>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0cb03cc0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0caf5898>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0ca3f4e0>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0ca09e10>, <tensorflow.contrib.tpu.python.tpu.keras_tpu_variables.ReplicatedVariable object at 0x7fec0c9cf860>]
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 11.36061954498291 secs
60000/60000 [==============================] - 39s 655us/sample - loss: 0.3495 - sparse_categorical_accuracy: 0.8974
Epoch 2/5
60000/60000 [==============================] - 4s 66us/sample - loss: 0.0626 - sparse_categorical_accuracy: 0.9814
Epoch 3/5
60000/60000 [==============================] - 4s 66us/sample - loss: 0.0338 - sparse_categorical_accuracy: 0.9900
Epoch 4/5
60000/60000 [==============================] - 4s 67us/sample - loss: 0.0203 - sparse_categorical_accuracy: 0.9941
Epoch 5/5
60000/60000 [==============================] - 4s 68us/sample - loss: 0.0110 - sparse_categorical_accuracy: 0.9972
INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(4,), dtype=tf.int32, name='core_id_30'), TensorSpec(shape=(4, 28, 28, 1), dtype=tf.float32, name='conv2d_1_input_10'), TensorSpec(shape=(4, 1), dtype=tf.float32, name='dense_3_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Cloning Adam {'lr': 0.0010000000474974513, 'beta_1': 0.8999999761581421, 'beta_2': 0.9990000128746033, 'decay': 0.0, 'epsilon': 1e-07, 'amsgrad': False}
INFO:tensorflow:Remapping placeholder for conv2d_1_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fec0b8f72e8> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 2.7844197750091553 secs
 9856/10000 [============================>.] - ETA: 0s - loss: 0.1032 - sparse_categorical_accuracy: 0.9698INFO:tensorflow:New input shapes; (re-)compiling: mode=eval (# of cores 8), [TensorSpec(shape=(2,), dtype=tf.int32, name='core_id_30'), TensorSpec(shape=(2, 28, 28, 1), dtype=tf.float32, name='conv2d_1_input_10'), TensorSpec(shape=(2, 1), dtype=tf.float32, name='dense_3_target_10')]
INFO:tensorflow:Overriding default placeholder.
INFO:tensorflow:Remapping placeholder for conv2d_1_input
INFO:tensorflow:KerasCrossShard: <tensorflow.python.keras.optimizers.Adam object at 0x7fec0b8f72e8> []
INFO:tensorflow:Started compiling
INFO:tensorflow:Finished compiling. Time elapsed: 1.8734314441680908 secs
10000/10000 [==============================] - 9s 921us/sample - loss: 0.1027 - sparse_categorical_accuracy: 0.9699
[0.1027357829653658, 0.96989995]

モデル変換→コンパイルの順序にすると少し速くなった。

Kerasでクラス分類モデルの出力をlogitsにする

Google ColabでTPUを使うには、今のところフレームワークにTesorFlow(Keras)を使う必要がある。
Kerasで将棋AI用のモデル定義を行っていて、ChainerではできてKerasでは簡単にできない問題にぶつかった。

Kerasでクラス分類のモデルを定義して学習する際、通常の方法では、出力層のSoftmax活性化関数を含めたモデルしか定義できない。

学習済みモデルを使って、ボルツマン分布で確率を出力したい場合には、出力として、Softmax活性化関数に入力する前のロジットを取得したい。
ChainerのようなDefine-by-Runのフレームワークでは途中の値を取り出すのは容易だが、Kerasはモデルを事前にコンパイルするため途中の値は取得できない。

Kerasでロジットを出力するモデルを定義して、できるだけ標準的な方法で学習できないか調べてみた。
python - Keras - how to get unnormalized logits instead of probabilities - Stack Overflow
このQAにヒントが書かれていた。

引数にfrom_logits=Trueを与えたsparse_categorical_crossentropyをラッピングしてカスタム損失関数を定義すればよい。

def sparse_categorical_crossentropy(y_true, y_pred):
    return tf.keras.backend.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)

metricsについては上記のQAには書かれていなかったが、試行錯誤の結果、以下のようにすればよいことがわかった。
metricsには、sparse_categorical_accuracyを使用したしたいが、モデルの出力がロジットになっているため、softmaxを計算した上で、sparse_categorical_accuracyに渡す必要がある。つまり、そのように処理を行うカスタム関数を定義すればよい。

def sparse_categorical_accuracy(y_true, y_pred):
    return tf.keras.metrics.sparse_categorical_accuracy(y_true, tf.nn.softmax(y_pred))

通常のモデルと、ロジットを出力するモデルのそれぞれの、モデル定義と学習方法は以下の通りになる。

通常のモデル

import tensorflow as tf
mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)

ロジットを出力するモデル

import tensorflow as tf

mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dense(10)
])

def sparse_categorical_crossentropy(y_true, y_pred):
    return tf.keras.backend.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)

def sparse_categorical_accuracy(y_true, y_pred):
    return tf.keras.metrics.sparse_categorical_accuracy(y_true, tf.nn.softmax(y_pred))

model.compile(optimizer='adam',
              loss=sparse_categorical_crossentropy,
              metrics=[sparse_categorical_accuracy])

model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)

高速なPythonの将棋ライブラリを作る

python-shogiは、Pythonで扱える非常に役立つ将棋ライブラリですが、速度が遅いのが用途によっては欠点になります。
公式サイトにも記述されていますが、速度よりもシンプルに抽象的に扱えることが目的となっています。

しかし、機械学習の用途に使用しようとする速度の遅さがネックになります。
そこでPythonからもできるだけ高速に動作する将棋ライブラリを作成することにしました。

python-shogiの内部では、盤面はビットボードで表現されていますが、Pythonのビット演算は非常に遅くボトルネックとなっています。
ビット演算部分をC++で開発して、Pythonから呼び出せるようにすることで速度の改善が見込まれます。
C++で将棋ライブラリを一から作成するのもロマンがありますが、ほとんど既存のライブラリをまねるだけになるため、C++部分にAperyのソースコードを使用させてもらい、AperyをPythonから呼び出せるようにする形で作成することにします。

C++Pythonのブリッジの方法には、Boost.PythonやPyBind11やcythonなどがありますが、cythonを使うことにしました。
理由は、cythonではインストール時に自動でビルドすることが可能で、Google Colabのようなクラウドでのインストールも容易になるためです。

C++Pythonバインディング

cythonでC++のクラスや関数を呼び出すことは比較的容易で、拡張子が、.pyxのファイルを作成して、Pythonとほぼ同じ構文で

cdef extern from "cshogi.h":
	cdef cppclass __Board:
		__Board() except +
		__Board(const string& sfen) except +
		bool set_hcp(const char* hcp)

のように、extern宣言を行うだけでPythonで定義したクラスとほぼ同様に扱うことができます。
型の変換も自動で行われるため、ほとんど気にすることはありません。
ただし、ポインタで配列を渡す場合などは、少しトリックが必要になります。

Pythonからポインタを渡すには、cythonのメモリビュー機能を使って、Numpyのndarrayからポインタに変換します。

def func(data)
    cdef char[::1] = data
    c_func(&data[0])

参考:Typed Memoryviews — Cython 3.0.0a10 documentation

Aperyから借用したソース

Aperyのソースコードから盤面管理と合法手生成の部分を残して、それ以外は削除しました。
PositionクラスはSearcherと密結合になっていましたが、Searcherが不要になるように修正を行いました。

Aperyにある局面を圧縮形式(HuffmanCodedPos)で保存、読み込みする機能も、Pythonから使用できるようにしています。
これにより、機械学習での局面の保存、読み込みが速くなります。

hcpでの保存、読み込みの例
import cshogi
import numpy as np
board = cshogi.Board()
# save
hcp = np.empty(1, dtype=cshogi.HuffmanCodedPos)
board.to_hcp(hcp[0]['hcp'])

# load
board.set_hcp(hcp[0]['hcp'])

独自に追加した機能

独自に追加した機能としては、python-shogiと同じようにCSAファイルの読み込みが行える機能を用意しました。

また、やねうら王の学習局面データ形式(PackedSfen)も読み込めるようにしました。

インストール

GitHubからcloneして、pipでインストールします。

git clone https://github.com/TadaoYamaoka/cshogi.git
cd cshogi
pip install -e .

直接GitHubからインストールすることもできます。

pip install git+https://github.com/TadaoYamaoka/cshogi --no-cache-dir

インストールパッケージ名前は、cshogiで、importする際のモジュール名も、cshogiになります。
※2019/2/28 インストールパッケージ名をモジュール名と合わせました。

使用例

盤を作成して、開始局面で合法手を生成して表示し、1手指す処理の例です。

import cshogi
board = cshogi.Board()
# legal moves
for move in board.legal_moves:
    print(cshogi.move_to_usi(move))
# move
move = board.push_usi('7g7f')
# print board
print(board)

注意点

文字列の扱い

文字列はすべてバイト列で扱います。
Python3では、文字列リテラルにbをプレフィックスとして付けます。

※2019/8/25 引数の文字列をバイト文字列から標準の文字列型に修正しました。

指し手の扱い

指し手は、数値で表します。
python-shogiでは、指し手はMoveクラスとして扱い、便利なメソッドが用意されていますが、速度重視でクラスにしていません。
その代わり、いくつかのヘルパーメソッドを用意しています。

  • move_to(move)
  • move_from(move)
  • move_is_promotion(move)
  • move_is_drop(move)
  • move_to_usi(move)
  • move_to_csa(move)

合法手チェック

速度を重視して着手を適用する際、合法手チェックを行っていません。
誤ったmoveを渡すと盤面データが壊れて、その後アクセス違反などでプログラムが異常終了する場合があります。

undo処理

python-shogiと異なり、局面を戻す際に、moveが必要になります。

# undo move
board.pop(move)
座標系

座標系は、Aperyに準拠しています。python-shogiとは異なるので注意が必要です。詳細は、こちらの記事を参照してください。

駒の扱い

駒に対応する数値もAperyに準拠しています。python-shogiとは異なるので注意が必要です。
詳細は、cshogi.pyxのPIECES、PIECE_TYPESの定義を確認してください。

速度比較

python-shogiと速度を比較しました。
比較には、機械学習を想定して、CSA形式の棋譜を読み込んで、棋譜を再生して、局面をリストに格納する処理を使用しました。
python-shogiでは、ビットボードをコピーしてリストに保存し、cshogiでは、hcpe形式でリストに格納します。
比較用コード:cshogi/read_csa.py at master · TadaoYamaoka/cshogi · GitHub

棋譜100ファイルを処理した時間の比較
python-shogi 0.6229 秒
cshogi 0.0629 秒

cshogiの方が、9.9倍高速です。

今後の予定

Google ColabでPythonオンリーでAlpha Zeroと同じ強化学習ができるようにすることを目標にしています。
そのために必要な機能拡充を行っていく予定です。

技術書典6に当選したので、強化学習をテーマに本を出す予定です。
本にするには、読者の実行環境を考慮する必要がありますが、誰でも同じように実行できるGoogle Colabが最適だと思います。
作ったのは、そのための準備でもあります。

今のところ使用方法とか説明がないので、使い方の詳細はソースを読む必要があります。
余裕ができたら使い方の説明も作成したいと思います。

補足

おそらくまだバグが残っています。
プルリクもお待ちしています。

github.com

2019/2/11 追記

Google Colabで使用する方法

現状PyPIに登録していないため、Google Colabで使用するには、GitHubからcloneしてローカルのファイルでインストールする必要がある。
そのためには、ノードブックで以下のコードを実行する。

!git clone https://github.com/TadaoYamaoka/cshogi.git
%cd cshogi
!pip install -e .
%cd ..

これで、cshogiをimportして使用できるようになる。

2019/2/16 追記

直接GitHubからインストールすることもできる。

!pip install git+https://github.com/TadaoYamaoka/cshogi --no-cache-dir
アンインストール

ランタイムをクリアすれば、インストールの状態もクリアされるが、個別にアンインストールしたい場合は、

!pip uninstall cshogi

を実行すればよい。ロード済みのモジュールはそのまま使用される。再インストールして再読み込みさせるにはランタイムの再起動が必要になる。

将棋AIの進捗 その26(自己対局による強化学習の経過2)

前回から時間が空いたが、自己対局による強化学習を続けている。

10ブロック、192フィルタのモデルの自己対局による学習が、79サイクル※回したところで飽和気味になったため、10ブロックのモデルからパラメータを転移して15ブロックのモデルで強化学習を行うことにした。
※1サイクルで、250万局面を生成(55サイクルまでは500万局面にしていた)

自己対局の速度比較

10ブロックの自己対局では43.1局面/秒で生成できていたが、15ブロックにしたことで37.30局面/秒になり、局面の生成速度は、79.6%に低下した。
250万局面の生成時間は、10ブロックは17時間12分、15ブロックは19時間56分となった。

モデル学習速度

10ブロックでは、250万局面×過去10サイクルの学習に、2時間45分だったところ、15ブロックでは、4時間4分となり、学習速度は、1.47倍になった。

自己対局とモデルの学習を合わせると、1サイクルあたり、

10ブロック 19時間57分
15ブロック 24時間42分

となった。

精度比較

70サイクルまで学習した10ブロックのモデルのパラメータを転移学習して、10ブロックのモデルで生成した71~80サイクルのデータを学習して、そこから自己対局を105サイクル(1サイクル目を81サイクルとしてカウント)まで行った結果、ほぼ飽和した。
floodgateのR3500以上の棋譜との一致率は以下のようになった。

モデル policy accuracy value accuracy
10ブロック 0.44010064 0.7129697
15ブロック 0.44107565 0.7058925

policyは一致率が高くなっているが、valueは大幅に低下してしまった。

バッチサイズの変更

valueの精度はバッチサイズを増やすと安定する傾向があるので、バッチサイズを256から1024に増やして、学習をやり直した。
99サイクルまで学習した結果、floodgateのR3500以上の棋譜との一致率は以下のようになった。

モデル policy accuracy value accuracy
10ブロック 0.44010064 0.7129697
15ブロック 0.44053894 0.7124748

policyの一致率はわずかに上がり、valueの一致率はほぼ同じになった。

強さの比較

10ブロックのモデルと、15ブロックのモデルで、同じ探索プログラムで1手3秒で100回対局を行い強さを比較した。
結果は、以下の通りとなった。

対局数100 先手勝ち47(47%) 後手勝ち51(51%) 引き分け2
selfplay079
勝ち62(62%) 先手勝ち30(30%) 後手勝ち32(32%)
selfplay099_b15
勝ち36(36%) 先手勝ち17(17%) 後手勝ち19(19%)

※selfplay079が10ブロック、selfplay099_b15が15ブロック
10ブロックのモデルの方が強いという結果となった。
15ブロックのモデルの精度の向上はわずかだが、探索速度は、初期局面で、

モデル プレイアウト/秒
10ブロック 5905
15ブロック 4214

となり、71.4%に低下している。
探索速度低下を補うには0.1%程度のモデルの精度向上では足りなかったようだ。

考察

モデルのブロック数を増やして、パラメータを転移学習する方法では、モデルの精度を上げることはできなかった。
転移は行わず初期値から生成済みの学習局面を学習させた方が良かったかもしれない。

15ブロックのモデルの検証にやり直しを含めて、2ヵ月かかっているので、再実験は一旦保留して、10ブロックのモデルの精度向上の方に取り組むことにする。
15ブロックのモデルで、バッチサイズを増やすと効果があったので、10ブロックのモデルでもバッチサイズを増やしてみる。

また、10ブロックのモデルと15ブロックのモデルを対局させた際、以下の局面で、詰み探索を入れているため先手が詰みを見つけている局面で後手は自分が優勢と考えていた。
f:id:TadaoYamaoka:20190204211758p:plain
先手側も詰み探索を除くと後手優勢と判断しており、終盤の詰みの読み漏れが起きていた。
詰みが絡む局面をうまく学習できていない可能性があるため、自己対局の対局打ち切りの閾値に問題があるのではないかと考えている。
自己対局のサイクルを早めるため、勝率の推定値が閾値を超えた(もしくは下回った)場合に対局を打ち切っているが、終盤の判断に誤りがある状態では、上記のように自分が詰んでいる局面で誤って自分の勝ちと判断して閾値を超えると、誤った学習を行い、それが強化されるという悪循環になる。

自己対局による強化学習では、学習の初期段階では自分自身の推定値で打ち切るのは問題がありそうである。
AlphaGo Zeroでは、一定数のゲームでは終局までプレイして勝率の推定値が正しいか検証を行っていた。
このプロセスは省いてはいけないのかもしれない。

対局打ち切り条件の変更

AlphaGo Zeroと同じように検証してもよいが、以前にDfPnによる高速な詰み探索を作成していたので、閾値による対局打ち切りを行わず、詰み探索で詰みが見つかるまで対局を行うように変更した。
GPUで計算中に詰み探索が完了するように、探索ノード数の上限を調整することで、自己対局による局面生成速度をほとんど落とすことなく詰み探索を組み込むことができた。
具体的な実装としては、1GPUにつき256エージェントの探索の推論要求バッチ処理しているため、詰み探索をエージェントにつき1スレッドで並列処理すると、以前に問題となった大群の問題が発生してしまうため、各エージェントの詰み探索の要求を1スレッドで順次処理するようにした。
エージェントごとに探索している局面と、何番目のプレイアウトを処理中かが異なるため、詰み探索を行うのは手番の切り替わりのタイミングのみなので、一度に詰み探索の要求が集中することは少ない。
そのため、1スレッドでもGPU計算中に詰み探索が完了するように調整できた。
探索ノード数を10万ノードに制限しているが、場合によっては20手以上の詰みも探索できていた。

これで、自己対局を回して、上記の局面のような誤った終盤の判断が改善するか様子を見たい。

余談

世界コンピュータ選手権申し込みました。
昨年よりは強くはなっていると思うので、なんとか1次予選は通過したいところです。