TPU 训练在训练过程中冻结答案

【问题标题】：TPU training freezes in the middle of trainingTPU 训练在训练过程中冻结
【发布时间】：2019-07-28 10:51:40
【问题描述】：

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

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)

【问题讨论】：

标签： neural-network cloud google-compute-engine tpu

【解决方案1】：

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

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

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

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

最后我建议你看看Training Mask RCNN on Cloud TPU 文档。

【讨论】：