使用Google colab在TPU上训练CNN

Question

我需要微调大型图像数据集上的CNN（vgg16）。我正在使用Google colab，并且需要使用TPU来加速培训。按照documentation example中的说明，我使用的是tf.keras而不是keras，在创建模型，替换完全连接的层并指定可训练的层之后，我使用{{1 }}损失函数，并创建了火车生成器：

categorical_crossentropy

然后我将模型转换为兼容的tensorflow TPU模型：

model.compile(loss='categorical_crossentropy', 
          optimizer= tensorflow.train.RMSPropOptimizer(learning_rate=1e-4),
          metrics = ['accuracy'])
train_data_gen = ImageDataGenerator(rescale=1./255,rotation_range = 20, 
                                width_shift_range = 0.2, 
                                height_shift_range = 0.2, 
                                horizontal_flip = True)
train_gen= train_data_gen.flow_from_directory('/content/drive/My Drive/data/train', 
                                          target_size=(224, 224), 
                                          batch_size = 1024, 
                                          class_mode='categorical' )

当我尝试使用以下命令训练模型时：

tf.logging.set_verbosity(tf.logging.INFO)
tpu_model = tf.contrib.tpu.keras_to_tpu_model(model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
    tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

我收到此错误消息：

history = tpu_model.fit_generator(train_gen, 
                    steps_per_epoch=train_gen.samples/train_gen.batch_size, 
                    epochs=1000, 
                    verbose=2, 
                    shuffle= False)

有人可以帮我解决这个问题吗？