概要
- Apple M1 に最適化されたTensorflow 1 を使って機械翻訳のチュートリアル2をやってみた。
- LSTMやDenseは最適化されているが、GRUやEmbeddingは最適化されていないっぽい。
- 最適化されている関数もパラメータに制約があることがある。
- 既存のコードに少し(?)手を入れると数倍速くなるので、M1のコスパはやはり良い。
はじめに
Apple M1 について、@tomoyaeibuさんの2017年のIntel Macと比べて4倍速いという報告3 を読みました。
僕も新しいMac miniを買っており、現実的な問題で何か試したくなったため、Tensorflowのチュートリアルにある機械翻訳タスク 2 にチャレンジしてみました。
タスクの詳細は脚注のリンク先を見て欲しいのですが、注意機構（アテンション）付きのエンコーダー・デコーダーモデルを作って、スペイン語から英語への機械翻訳をやろう、というものです。
チュートリアルのコードをそのまま実装すると、最適化されたTensorflowでも、ものすごく時間がかかります。
本稿では、この遅さの原因を調べて、とりあえずの解決をやってみた経験から、
Apple M1において自然言語処理の具体的なタスクを実装する場合の注意点、特にLSTMの使い方について書きます。
最適化されたTensorflowのインストールなんかは省略しておりますので、脚注から他の方の記事をお読みください。
僕が使ったTensorflowのバージョンは、実験時点の最新である tensorflow-macos 0.1a1 でした。
詳細な環境情報は本稿の後ろの方にまとめてあります。
GRUは最適化されていない
チュートリアル本文中には「このサンプルは P100 GPU 1基で実行した場合に約 10 分かかります。」とありますが、同じコードを新しいMac miniで試すと、1エポック550secかかり、100分くらいはかかってしまいそうです。
（本稿末尾の速度比較の節を参照してください。）
もちろんGPUはとんでもなく速いので、10倍くらいはしょうがないという気もしますが、せっかくの新しいCPUなのだからもう少し頑張って欲しいところです。
Apple M1 に最適化されたTensorflowである、tensorflow-macosは、内部的にML Compute4というフレームワークを使っています。
ML Computeが呼べる関数には、再帰ニューラルネットワーク用のレイヤとして LSTM（長期短期記憶）は入っているのですが、今回のチュートリアルで使われているGRU（ゲート付き回帰ユニット）は入っていません。
このため、tensorflow-macosにおいても、LSTMにはML Computeを使うためのコードが書かれている一方で、GRUについては特に手が入っていないようです。
そこで、エンコーダ、デコーダのそれぞれの再帰ニューラルネットワーク層の実装を、GRUからLSTMに変更してみることにしました。
（余談） ML Computeが使えるレイヤの調べ方
tensorflow-macosをインストールした際に作ったvenv環境の下にあるlibディレクトリに対して、APPLE_MLCOMPUTEという文字列で grep をかけると、ML Computeが使えるtensorflowのメソッドがざっくり分かりそうです。
GRUをLSTMモドキに変更
いきなりですが、コードを示します。
元のGRUを使った部分はコメントアウトして残してあり、似たような機能を持たせた層をLSTMで実装しています。
（実行可能なコードの全体は、本稿末尾にあります。）
エンコーダ部分（LSTMモドキ版）
class Encoder(tf.keras.Model):
def __init__(self, vocab_size, embedding_dim, enc_units, batch_sz):
super(Encoder, self).__init__()
self.batch_sz = batch_sz
self.enc_units = enc_units
self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
# 元のコードではGRUで実装されている
# self.gru = tf.keras.layers.GRU(self.enc_units,
# return_sequences=True,
# return_state=True,
# recurrent_initializer='glorot_uniform')
# LSTMモドキ で書き直す
self.lstm = tf.keras.layers.LSTM(self.enc_units,
return_sequences=True,
return_state=False, # TrueだとML Computeが使えない
recurrent_initializer='glorot_uniform')
def call(self, x, hidden):
x = self.embedding(x)
# output, state = self.gru(x, initial_state=hidden)
# 現状LSTMの隠れ状態ベクトルを取り出す方法がないので
# 最後の出力ベクトルを無理矢理次の隠れ状態ベクトルにコピーして使う。
output = self.lstm(x, initial_state=[hidden, hidden])
state = output[:, -1, :]
return output, state
def initialize_hidden_state(self):
return tf.zeros((self.batch_sz, self.enc_units))
ML Compute を使う際の制約
さて、上記のコードはぱっと見では GRUという部分をLSTMに変更しただけに見えます。
しかしよく見ると、再帰ネットワークの最終状態ベクトルを返すかどうかを選択するreturn_stateオプションについて、オリジナルのGRU版ではTrueである一方、僕のLSTM版ではFalseにしています。
実は現在のtensorflow-macosでは、このオプションがTrueの場合は、ML Computeを使わないことにしているようです。
LSTMの内部状態をPythonとML Computeの間でやりとりするのが大変だからだろうなと想像しています。
GRUは隠れ状態ベクトルが一つだけで、その最終状態は出力されるベクトルシーケンスの最後のものと一致するのですが、LSTMは隠れ状態を2つ持つうえ、最終状態は出力ベクトルシーケンスとは違うものです。
…違うものなのですが、今回の機械翻訳という目的にとっては、この辺りの細かいことは、注意機構辺りがうまいことやってくれるやろうといい加減に考えて、出力ベクトルシーケンスを単に複製して、続く隠れ状態ベクトルの初期値としています。
学習を進めた際の損失の減り方なんかを見ると、どうも大丈夫そうでした。
エンコーダに対して行ったのと同様の改変をデコーダについても行う必要があります。
速度比較
では、エンコーダ部分とデコーダ部分に上記の改変を加えた上で、実際の学習速度の比較を行った結果を書きます。
実行環境
以下の3つの環境で比較しました。
新旧プロセッサの比較なのですが、参考までにGoogle ColabのGPU環境を使った場合も載せています。
| Mac mini | 旧MacBook Pro | （参考）Google Colab | |
|---|---|---|---|
| Model | (M1, 2020) | (Retina, 13-inch, Early 2015) | - |
| OS | macOS Big Sur | macOS Catalina | - |
| CPU | M1 | Dual-Core Intel Corei5 2.7GHz | - |
| GPU | M1 | Intel Iris Plus Graphics 6100 1536MB | Tesla T4 14.7GiB |
| RAM | 16GB | 16GB | - |
| Python | python 3.8.2 | python 3.8.6 (Homebrew) | - |
| Tensorflow | tensorflow-macos 0.1a1 | tensorflow 2.3.1 (Homebrew) | - |
結果
| #1 | #2 | #3 | （参考） #4 | |
|---|---|---|---|---|
| 環境 | Mac mini | Mac mini | 旧MacBook Pro | Google Colab |
| コード | LSTMモドキ版 | オリジナルGRU版 | LSTMモドキ版 | LSTMモドキ版 |
| 1エポックにかかった時間 | 162〜167sec | 537〜557sec | 888〜902sec | 36〜37sec |
まとめ
GRUが使えないこと、LSTMを使う際にもちょっとした制約があることなど、Apple M1で自然言語処理を効率的に動かすには細かい調整が必要でこれまで使っていたコードをそのままとはいかず、本格的に使うにはまだちょっと手を出しづらい感じですね。
ただし、新しい Apple M1 の実力は、最適化が効かない状態でも、5年前のIntel Mac搭載のMBPの2倍近く高速で、調整を行うことで、さらに3倍くらいは速くできそうです。
Google Colabで使える速いGPUには全く及ばないのですが、10万円前後で買える機械としてはまずまずの性能といえそうです。
荒削りのライブラリを読み解いていくのが楽しいような人にとっては、Google ColabのGPU無料枠を使い切った時間なんかに良いのではないでしょうか(^^;
（付録） ソースコード全体
今回の実験を行った、LSTMモドキ版の実行可能なソース全体を載せています。
# コードの大部分は以下のチュートリアルによる。
# https://www.tensorflow.org/tutorials/text/nmt_with_attention?hl=ja
# 元のコンテンツは Creative Commons Attribution 4.0 License の下で公開されている。
import io
import os
import re
import time
import unicodedata
import tensorflow as tf
path_to_zip = tf.keras.utils.get_file(
'spa-eng.zip', origin='http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip',
extract=True)
path_to_file = os.path.dirname(path_to_zip) + "/spa-eng/spa.txt"
# ユニコードファイルを ascii に変換
def unicode_to_ascii(s):
return ''.join(c for c in unicodedata.normalize('NFD', s)
if unicodedata.category(c) != 'Mn')
def preprocess_sentence(w):
w = unicode_to_ascii(w.lower().strip())
# 単語とそのあとの句読点の間にスペースを挿入
# 例： "he is a boy." => "he is a boy ."
# 参照：- https://stackoverflow.com/questions/3645931/python-padding-punctuation-with-white-spaces-keeping-punctuation
w = re.sub(r"([?.!,¿])", r" \1 ", w)
w = re.sub(r'[" "]+', " ", w)
# (a-z, A-Z, ".", "?", "!", ",") 以外の全ての文字をスペースに置き換え
w = re.sub(r"[^a-zA-Z?.!,¿]+", " ", w)
w = w.rstrip().strip()
# 文の開始と終了のトークンを付加
# モデルが予測をいつ開始し、いつ終了すれば良いかを知らせるため
w = '<start> ' + w + ' <end>'
return w
en_sentence = u"May I borrow this book?"
sp_sentence = u"¿Puedo tomar prestado este libro?"
print(preprocess_sentence(en_sentence))
print(preprocess_sentence(sp_sentence).encode('utf-8'))
# 1. アクセント記号を除去
# 2. 文をクリーニング
# 3. [ENGLISH, SPANISH] の形で単語のペアを返す
def create_dataset(path, num_examples):
lines = io.open(path, encoding='UTF-8').read().strip().split('\n')
word_pairs = [[preprocess_sentence(w) for w in l.split('\t')] for l in lines[:num_examples]]
return zip(*word_pairs)
def max_length(tensor):
return max(len(t) for t in tensor)
def load_dataset(path, num_examples=None):
# クリーニングされた入力と出力のペアを生成
targ_lang, inp_lang = create_dataset(path, num_examples)
input_tensor, inp_lang_tokenizer = tokenize(inp_lang)
target_tensor, targ_lang_tokenizer = tokenize(targ_lang)
return input_tensor, target_tensor, inp_lang_tokenizer, targ_lang_tokenizer
def tokenize(lang):
lang_tokenizer = tf.keras.preprocessing.text.Tokenizer(
filters='')
lang_tokenizer.fit_on_texts(lang)
tensor = lang_tokenizer.texts_to_sequences(lang)
tensor = tf.keras.preprocessing.sequence.pad_sequences(tensor,
padding='post')
return tensor, lang_tokenizer
# このサイズのデータセットで実験
num_examples = 30000
input_tensor, target_tensor, inp_lang, targ_lang = load_dataset(path_to_file, num_examples)
# ターゲットテンソルの最大長を計算
max_length_targ, max_length_inp = max_length(target_tensor), max_length(input_tensor)
# 80-20で分割を行い、訓練用と検証用のデータセットを作成
# sklearnが使えないので、自分で実装する
def train_test_split(it, tt, test_size=0.2):
l_val = int(len(it) * test_size)
l_train = len(it) - l_val
l_val_target = int(len(tt) * test_size)
l_train_target = len(tt) - l_val_target
return it[:l_train], it[l_train:], tt[:l_train_target], tt[l_train_target:]
input_tensor_train, input_tensor_val, target_tensor_train, target_tensor_val = \
train_test_split(input_tensor, target_tensor, test_size=0.2)
# 長さを表示
print(len(input_tensor_train), len(target_tensor_train), len(input_tensor_val), len(target_tensor_val))
def convert(lang, tensor):
for t in tensor:
if t != 0:
print("%d ----> %s" % (t, lang.index_word[t]))
print("Input Language; index to word mapping")
convert(inp_lang, input_tensor_train[1])
print("Target Language; index to word mapping")
convert(targ_lang, target_tensor_train[1])
BUFFER_SIZE = len(input_tensor_train)
BATCH_SIZE = 64
steps_per_epoch = len(input_tensor_train) // BATCH_SIZE
embedding_dim = 256
units = 1024
vocab_inp_size = len(inp_lang.word_index) + 1
vocab_tar_size = len(targ_lang.word_index) + 1
dataset = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train)).shuffle(BUFFER_SIZE)
dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)
example_input_batch, example_target_batch = next(iter(dataset))
print(example_input_batch.shape, example_target_batch.shape)
class Encoder(tf.keras.Model):
def __init__(self, vocab_size, embedding_dim, enc_units, batch_sz):
super(Encoder, self).__init__()
self.batch_sz = batch_sz
self.enc_units = enc_units
self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
# self.gru = tf.keras.layers.GRU(self.enc_units,
# return_sequences=True,
# return_state=True,
# recurrent_initializer='glorot_uniform')
self.lstm = tf.keras.layers.LSTM(self.enc_units,
return_sequences=True,
return_state=False,
recurrent_initializer='glorot_uniform')
def call(self, x, hidden):
x = self.embedding(x)
# output, state = self.gru(x, initial_state=hidden)
output = self.lstm(x, initial_state=[hidden, hidden])
state = output[:, -1, :]
return output, state
def initialize_hidden_state(self):
return tf.zeros((self.batch_sz, self.enc_units))
encoder = Encoder(vocab_inp_size, embedding_dim, units, BATCH_SIZE)
# サンプル入力
sample_hidden = encoder.initialize_hidden_state()
sample_output, sample_hidden = encoder(example_input_batch, sample_hidden)
print('Encoder output shape: (batch size, sequence length, units) {}'.format(sample_output.shape))
print('Encoder Hidden state shape: (batch size, units) {}'.format(sample_hidden.shape))
class BahdanauAttention(tf.keras.layers.Layer):
def __init__(self, units):
super(BahdanauAttention, self).__init__()
self.W1 = tf.keras.layers.Dense(units)
self.W2 = tf.keras.layers.Dense(units)
self.V = tf.keras.layers.Dense(1)
def call(self, query, values):
# hidden shape == (batch_size, hidden size)
# hidden_with_time_axis shape == (batch_size, 1, hidden size)
# スコアを計算するためにこのように加算を実行する
hidden_with_time_axis = tf.expand_dims(query, 1)
# score shape == (batch_size, max_length, 1)
# スコアを self.V に適用するために最後の軸は 1 となる
# self.V に適用する前のテンソルの shape は (batch_size, max_length, units)
score = self.V(tf.nn.tanh(
self.W1(values) + self.W2(hidden_with_time_axis)))
# attention_weights の shape == (batch_size, max_length, 1)
attention_weights = tf.nn.softmax(score, axis=1)
# context_vector の合計後の shape == (batch_size, hidden_size)
context_vector = attention_weights * values
context_vector = tf.reduce_sum(context_vector, axis=1)
return context_vector, attention_weights
attention_layer = BahdanauAttention(10)
attention_result, attention_weights = attention_layer(sample_hidden, sample_output)
print("Attention result shape: (batch size, units) {}".format(attention_result.shape))
print("Attention weights shape: (batch_size, sequence_length, 1) {}".format(attention_weights.shape))
class Decoder(tf.keras.Model):
def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz):
super(Decoder, self).__init__()
self.batch_sz = batch_sz
self.dec_units = dec_units
self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
# self.gru = tf.keras.layers.GRU(self.dec_units,
# return_sequences=True,
# return_state=True,
# recurrent_initializer='glorot_uniform')
self.lstm = tf.keras.layers.LSTM(self.dec_units,
return_sequences=True,
return_state=False,
recurrent_initializer='glorot_uniform')
self.fc = tf.keras.layers.Dense(vocab_size)
# アテンションのため
self.attention = BahdanauAttention(self.dec_units)
def call(self, x, hidden, enc_output):
# enc_output の shape == (batch_size, max_length, hidden_size)
context_vector, attention_weights = self.attention(hidden, enc_output)
# 埋め込み層を通過したあとの x の shape == (batch_size, 1, embedding_dim)
x = self.embedding(x)
# 結合後の x の shape == (batch_size, 1, embedding_dim + hidden_size)
x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)
# 結合したベクトルを GRU 層に渡す
# output, state = self.gru(x)
output = self.lstm(x)
state = output[:, -1, :]
# output shape == (batch_size * 1, hidden_size)
output = tf.reshape(output, (-1, output.shape[2]))
# output shape == (batch_size, vocab)
x = self.fc(output)
return x, state, attention_weights
decoder = Decoder(vocab_tar_size, embedding_dim, units, BATCH_SIZE)
sample_decoder_output, _, _ = decoder(tf.random.uniform((64, 1)),
sample_hidden, sample_output)
print('Decoder output shape: (batch_size, vocab size) {}'.format(sample_decoder_output.shape))
optimizer = tf.keras.optimizers.Adam()
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction='none')
def loss_function(real, pred):
mask = tf.math.logical_not(tf.math.equal(real, 0))
loss_ = loss_object(real, pred)
mask = tf.cast(mask, dtype=loss_.dtype)
loss_ *= mask
return tf.reduce_mean(loss_)
checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(optimizer=optimizer,
encoder=encoder,
decoder=decoder)
@tf.function
def train_step(inp, targ, enc_hidden):
loss = 0
with tf.GradientTape() as tape:
enc_output, enc_hidden = encoder(inp, enc_hidden)
dec_hidden = enc_hidden
dec_input = tf.expand_dims([targ_lang.word_index['<start>']] * BATCH_SIZE, 1)
# Teacher Forcing - 正解値を次の入力として供給
for t in range(1, targ.shape[1]):
# passing enc_output to the decoder
predictions, dec_hidden, _ = decoder(dec_input, dec_hidden, enc_output)
loss += loss_function(targ[:, t], predictions)
# Teacher Forcing を使用
dec_input = tf.expand_dims(targ[:, t], 1)
batch_loss = (loss / int(targ.shape[1]))
variables = encoder.trainable_variables + decoder.trainable_variables
gradients = tape.gradient(loss, variables)
optimizer.apply_gradients(zip(gradients, variables))
return batch_loss
EPOCHS = 10
for epoch in range(EPOCHS):
start = time.time()
enc_hidden = encoder.initialize_hidden_state()
total_loss = 0
for (batch, (inp, targ)) in enumerate(dataset.take(steps_per_epoch)):
batch_loss = train_step(inp, targ, enc_hidden)
total_loss += batch_loss
if batch % 100 == 0:
print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,
batch,
batch_loss.numpy()))
# 2 エポックごとにモデル（のチェックポイント）を保存
if (epoch + 1) % 2 == 0:
checkpoint.save(file_prefix=checkpoint_prefix)
print('Epoch {} Loss {:.4f}'.format(epoch + 1,
total_loss / steps_per_epoch))
print('Time taken for 1 epoch {} sec\n'.format(time.time() - start))
def evaluate(sentence):
sentence = preprocess_sentence(sentence)
inputs = [inp_lang.word_index[i] for i in sentence.split(' ')]
inputs = tf.keras.preprocessing.sequence.pad_sequences([inputs],
maxlen=max_length_inp,
padding='post')
inputs = tf.convert_to_tensor(inputs)
result = ''
hidden = [tf.zeros((1, units))]
enc_out, enc_hidden = encoder(inputs, hidden)
dec_hidden = enc_hidden
dec_input = tf.expand_dims([targ_lang.word_index['<start>']], 0)
for t in range(max_length_targ):
predictions, dec_hidden, attention_weights = decoder(dec_input,
dec_hidden,
enc_out)
predicted_id = tf.argmax(predictions[0]).numpy()
result += targ_lang.index_word[predicted_id] + ' '
if targ_lang.index_word[predicted_id] == '<end>':
return result, sentence
# 予測された ID がモデルに戻される
dec_input = tf.expand_dims([predicted_id], 0)
return result, sentence
def translate(sentence):
result, sentence = evaluate(sentence)
print('Input: %s' % (sentence))
print('Predicted translation: {}'.format(result))
# checkpoint_dir の中の最後のチェックポイントを復元
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
translate(u'hace mucho frio aqui.')
translate(u'esta es mi vida.')
translate(u'¿todavia estan en casa?')
# 翻訳あやまりの例
translate(u'trata de averiguarlo.')
-
"Neural machine translation with attention", TensorFlow Core > Tutorials. ↩
-
@tomoyaeibu 「Apple Silicon M1 でtensorflow-macosを実行したらめちゃくちゃ速かった。」 ↩
-
"ML Compute", Apple Developer ↩