如何在 TPU 中使用自定义数据集生成器？

Question

我的数据集很大（大约 13gb）。我有一个数据集的 hdf5 文件，我正在使用自定义生成器从数据集中加载批次。我的模型在 Kaggle GPU 上运行良好，但是当我切换到 TPU 时出现错误。下面是我的生成器函数和运行 model.fit 时收到的错误。

def generate_data():

while True: # Loop forever so the generator never terminates
    
    for _ in range(0, num_samples, BATCH_SIZE):
        # Get the samples you'll use in this batch
        offset=np.random.randint(num_samples)
        X_train = hdf5_file['data'][offset]
        X_train=X_train.transpose(1,2,0)
        X_train=X_train.astype(np.float32)
        X_train=(X_train-127.5)/127.5
        X_train = cv2.resize(X_train, dsize=(IMG_SHAPE[1],IMG_SHAPE[1]), interpolation=cv2.INTER_CUBIC)
        X_train = np.array(X_train)
        
        #yield the next training batch
        yield X_train

这是从生成器 fn 生成 tf 数据集的代码。

dataset = tf.data.Dataset.from_generator(generate_data, (tf.float32))
dataset = dataset.batch(BATCH_SIZE,drop_remainder=True)

现在这是我在使用上述数据集运行 model.fit 时收到的错误。

TypeError: in user code:

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py:571 train_function  *
    outputs = self.distribute_strategy.run(
/opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/tpu_strategy.py:174 run  **
    return self.extended.tpu_run(fn, args, kwargs, options)
/opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/tpu_strategy.py:867 tpu_run
    return func(args, kwargs)
/opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/tpu_strategy.py:916 tpu_function
    maximum_shape = tensor_shape.TensorShape([None] * rank)

TypeError: can't multiply sequence by non-int of type 'NoneType'

正如我提到的，代码在 GPU 上运行良好，无需任何修改。我应该怎么做才能让它在 TPU 上工作？

Answer 1

您是否尝试过在from_generator() 函数中定义输出的形状？