TPUでConv1Dのdilation_rateを指定するとReshape’s input dynamic dimension is decomposed into multiple output dynamic dimensions

ColabのTPUでdilation convolutionを使った時に表題のエラー。

tf.keras.layers.Conv1D(32, kernel_size=3, dilation_rate=2)

Reshape's input dynamic dimension is decomposed into multiple output dynamic dimensions, but the constraint i s ambiguous and XLA can't infer the output dimension

GPUでは問題なく動いていたので、TPU固有の問題なのだろうか？

根本的な解決ではないが、tf.dataでdrop_remainder=Trueにすると回避出来た。

dataset = tf.data.Dataset.from_tensor_slices(data).shuffle(1024).repeat()
dataset = dataset.batch(batch_size, drop_remainder=True)

データセット数を割り切れるバッチサイズならエラーは回避出来るのか？と思い、バッチサイズをデータセット数の約数にしてみたが、やはり表題のエラーは発生した。

drop_remainderを設定した場合、データセット数がバッチサイズで割り切れない場合に余りが切り捨てられてしまう。

全てのデータセットを1epochで学習させたい場合はデータセット数を割り切れるバッチサイズに設定するよう気をつける。

参考

InvalidArgumentError with `Conv1D` which has the same `dilation_rate` as the batch size, using `tf.keras.Model.fit`, on TPU · Issue #43690 · tensorflow/tensorflow

System information Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes O...

tf.data.Dataset | TensorFlow v2.16.1

Represents a potentially large set of elements.