Getting ValueError: input tensor Tensor to TPUStrategy.run() has unknown rank, 这在 Google Collab 中是不允许的

Question

我正在尝试在 Google Colab 中使用 TPU 训练 RNN 生成模型。您可以在 here 找到笔记本的完整代码。简而言之，我获取文本文件，根据序列和目标将它们切碎，然后从列表中创建一个 tf.data.Dataset。然后准备词汇表，并创建带有预设词汇表的 keras.TextVectorization 对象。然后准备一个 one-hot 数据集，它应该将序列（特征）返回为 (60, 107) 张量，将目标返回为 (107,) 张量。然后我创建一个简单的模型，里面有一个 LTSM 层

with strategy.scope():

并尝试循环训练模型。像这样的东西：

import tensorflow as tf
tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
strategy = tf.distribute.TPUStrategy(tpu)

#some code to make dataset (skipped)... and then

one_hot_dataset = dataset.map(lambda x, y: (tf.one_hot(text_vectorizer(x), 
                                                       depth=vocab_size, 
                                                       dtype='float32'), 
                                            tf.squeeze(tf.one_hot(text_vectorizer(y), 
                                                                  depth=vocab_size, 
                                                                  dtype='float32', 
                                                                  axis=1))))

batch_size = 1024 # I hope this can load a TPU sufficiently

one_hot_dataset = one_hot_dataset.batch(batch_size=batch_size,
                                        num_parallel_calls=4)

one_hot_dataset = one_hot_dataset.prefetch(buffer_size=tf.data.AUTOTUNE)

# making a model

import keras
from keras import layers, Model

with strategy.scope():
    inputs = keras.Input(shape=(maxlen, vocab_size), dtype='float32')
    lstm_output = layers.LSTM(128)(inputs)
    output = layers.Dense(vocab_size, activation='softmax')(lstm_output)

    model = Model(inputs, output)

    model.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.01),
                  loss='categorical_crossentropy')

# and finally

model.fit(one_hot_dataset, epochs=1)

然后我得到一个错误：

ValueError: in user code:

    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1284, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1268, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))

    ValueError: input tensor Tensor("cond/Identity_8:0", dtype=float32) to TPUStrategy.run() has unknown rank, which is not allowed

我什至不能谷歌。

模型在 GPU 上训练正常（但速度太慢）。在类似的情况下，我已经成功地使用锅炉代码与 TPU 一起飞行。我对我的数据集有点怀疑，但不幸的是无法理解这个问题。

Answer 1

你能设置共享数据集元素的形状吗（批量维度除外）并且打印

one_hot_dataset.element_spec

是正确的。

例如

def map_fn(x, y):
  first = tf.one_hot(text_vectorizer(x), depth=vocab_size, dtype='float32')
  first.set_shape((60, 107))
  second = tf.squeeze(tf.one_hot(text_vectorizer(y), depth=vocab_size,
                                 dtype='float32', axis=1))
  second.set_shape(my_shape)
  return (first, second)

one_hot_dataset = dataset.map(lambda x, y: map_fn(x, y))