我正在尝试使用Google Colab中的TPU训练RNN生成模型。笔记本的完整代码你可以找到here.简而言之，我取文本文件，将它们按序列和目标分割，然后制作一个tf。data。列表中的数据集。然后准备词汇，创建Keras。具有预设词汇表的TextVectorization对象。然后准备一个one-hot数据集，该数据集应该返回序列（特征）作为（60，107）Tensor，目标作为（107，）Tensor。然后创建一个简单的模型，其中包含一个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早些时候在类似的情况下成功。我有点怀疑我的数据集，但不幸的是不能理解这个问题。