tensorflowのoptimizerに関するエラーを解決したい

実現したいこと

transformer encoderを用いたモデルの学習を行いたい

前提

tensorflowを用いて, とあるモデルの学習を行おうとしています.
学習を実行しようとしたところ, 該当箇所でエラーとなっており自分なりに調べましたが解決に至りませんでした.
以下のサイトを参考に実装しています.
https://www.tensorflow.org/tutorials/text/transformer?hl=ja#%E3%82%AA%E3%83%97%E3%83%86%E3%82%A3%E3%83%9E%E3%82%A4%E3%82%B6%E3%83%BC

よろしくお願いいたします.

発生している問題・エラーメッセージ

Traceback (most recent call last):
  File "(省略)/train.py", line 85, in <module>
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "(省略)/miniconda3/envs/test/lib/python3.11/site-packages/keras/optimizers/adam.py", line 116, in __init__
    self._learning_rate = self._build_learning_rate(learning_rate)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "(省略)/miniconda3/envs/test/lib/python3.11/site-packages/keras/optimizers/optimizer.py", line 384, in _build_learning_rate
    learning_rate(self.iterations)
  File "(省略)/adam.py", line 17, in __call__
    arg1 = tf.math.rsqrt(step)
           ^^^^^^^^^^^^^^^^^^^
  File "(省略)/miniconda3/envs/test/lib/python3.11/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "(省略)/miniconda3/envs/test/lib/python3.11/site-packages/tensorflow/python/framework/ops.py", line 7262, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: Value for attr 'T' of int64 is not in the list of allowed values: bfloat16, half, float, double, complex64, complex128
        ; NodeDef: {{node Rsqrt}}; Op<name=Rsqrt; signature=x:T -> y:T; attr=T:type,allowed=[DT_BFLOAT16, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, DT_COMPLEX128]> [Op:Rsqrt] name:

該当のソースコード

train.py
1learning_rate = CustomSchedule(d_model)
2print(f'learning_rate: {learning_rate}')
3optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)
4print('optimizer is ready')

adam.py
1class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
2  def __init__(self, d_model, warmup_steps=4000):
3    super(CustomSchedule, self).__init__()
4
5    self.d_model = d_model
6    self.d_model = tf.cast(self.d_model, tf.float32)
7
8    self.warmup_steps = warmup_steps
9
10  def __call__(self, step):
11    arg1 = tf.math.rsqrt(step)
12    arg2 = step * (self.warmup_steps ** -1.5)
13
14    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)

試したこと

train.pyの3行目を以下のように変更してみたところ, optimizer is readyと出力されました.
おそらくlearning_rateの部分に問題があるのではないかと考えています.

python
1optimizer = tf.keras.optimizers.Adam(0.001, beta_1=0.9, beta_2=0.98, epsilon=1e-9)

meg_

2023/08/22 10:33

「発生している問題・エラーメッセージ」は掲載のもので全てですか？「d_model」とは何でしょうか？

48270-50070

2023/08/23 04:03

エラーメッセージについて, 全文を掲載しました. ファイル構成は一部省略しております. d_modelはモデルへの入力ベクトルの次元数となっております.

行動規範の内容に同意します

回答1件

ベストアンサー

下記コードのarg1 = tf.math.rsqrt(step)の前にstep = tf.cast(step, tf.float32)を追加しましょう。

def __call__(self, step):
    arg1 = tf.math.rsqrt(step)
    arg2 = step * (self.warmup_steps ** -1.5)

    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)

tf.math.rsqrt

Google Colaboratory(cpu)で実際に本日試したコードを載せます。

Python
1import tensorflow as tf
2
3class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
4  def __init__(self, d_model, warmup_steps=4000):
5    super(CustomSchedule, self).__init__()
6
7    self.d_model = tf.cast(d_model, tf.float32)
8    self.warmup_steps = warmup_steps
9
10  def __call__(self, step):
11    step = tf.cast(step, tf.float32)
12    arg1 = tf.math.rsqrt(step)
13    arg2 = step * (self.warmup_steps ** -1.5)
14
15    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
16
17learning_rate = CustomSchedule(128)
18print(f'learning_rate: {learning_rate}')
19optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)
20print('optimizer is ready')
21
22# learning_rate: <__main__.CustomSchedule object at 0x7bfec8bdb130>
23# optimizer is ready