TPU训练在训练过程中冻结

Question

我正在尝试使用TPU v3-8 1.12实例在TF 1.12中训练CNN回归网络。该模型使用XLA成功编译，开始了训练过程，但在1t纪元的半个迭代之后的一些地方冻结，并且什么也不做。我找不到问题的根源。

def read_tfrecord(example):
    features = {
        'image': tf.FixedLenFeature([], tf.string),
        'labels': tf.FixedLenFeature([], tf.string)
    }
    sample=tf.parse_single_example(example, features)
    image = tf.image.decode_jpeg(sample['image'], channels=3)
    image = tf.reshape(image, tf.stack([540, 540, 3]))
    image = augmentation(image)
    labels = tf.decode_raw(sample['labels'], tf.float64)
    labels = tf.reshape(labels, tf.stack([2,2,45]))
    labels = tf.cast(labels, tf.float32)
    return image, labels

def load_dataset(filenames):
    files = tf.data.Dataset.list_files(filenames)
    dataset = files.apply(tf.data.experimental.parallel_interleave(tf.data.TFRecordDataset, cycle_length=4))
    dataset = dataset.apply(tf.data.experimental.map_and_batch(map_func=read_tfrecord, batch_size=BATCH_SIZE, drop_remainder=True))
    dataset = dataset.apply(tf.data.experimental.shuffle_and_repeat(1024, -1))
    dataset = dataset.prefetch(buffer_size=1024)
    return dataset

def augmentation(img):
    image = tf.cast(img, tf.float32)/255.0
    image = tf.image.random_brightness(image, max_delta=25/255)
    image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
    image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
    image = tf.image.per_image_standardization(image)
    return image

def get_batched_dataset(filenames):
    dataset = load_dataset(filenames)
    return dataset


def get_training_dataset():
    return get_batched_dataset(training_filenames)

def get_validation_dataset():
    return get_batched_dataset(validation_filenames)

Answer 1

最有可能的原因是数据预处理功能中的问题，请查看故障排除文档Errors in the middle of training，它可能有助于获得指导。

我没有发现你的代码有什么奇怪的地方。

您是否正在使用Cloud Storage Buckets处理这些图像和文件？如果是，这些存储桶是否在同一区域？

您可以使用Cloud TPU Audit Logs来确定问题是否与系统中的资源或您访问数据的方式相关。

最后，我建议您查看一下Training Mask RCNN on Cloud TPU文档。