TensorFlow 2-tf.keras：如何使用 tf.data API 和 TFRecords 训练像 MTCNN 这样的 tf.keras 多任务网络

Question

最近我正在尝试使用 TFRecords 来训练 tf.keras 模型。由于对于 TensorFlow 2，最有效的方法是使用 tf.data API，我尝试用它来训练我的 MTCNN Keras 模型。但这让我感到困惑：

根据原始论文，不同的样本（pos、neg、part-face、landmark）参与不同部分的训练。并且每种样本在每个 mini-batch 中都有一个特定的比例，即对于一个 mini-batch，pos、neg、part-face 和 landmark 样本的比例应该是 1:3:1:2。

因此，例如，在我制作了四个 TFRecords 样本之后，我需要从 pos tfrecords 中提取 128 个样本，从 neg 中提取 384 个样本，从 part-face 中提取 128 个样本，从 landmark tfrecords 中提取 256 个样本并将它们混合成一个小批量。然后我需要在训练前洗牌。

我真的不知道在使用 TFRecords 和 tf.data API 时如何进行此操作。现在我只能通过逐批读取图像和标签来实现所有这些步骤，但它对于训练来说太慢了。所以我想知道是否有任何有效的方法可以做到这一点。

如有任何建议，我们将不胜感激。

2020.04.04 更新15:38
感谢@AAudibert，我认为 his/her 答案非常有效，我也想出了一种实现方法。以下是参考代码：

raw_pos_dataset = tf.data.TFRecordDataset(POS_TFRECORDS_PATH_LIST)
raw_neg_dataset = tf.data.TFRecordDataset(NEG_TFRECORDS_PATH_LIST)
raw_part_dataset = tf.data.TFRecordDataset(PART_TFRECORDS_PATH_LIST)
raw_landmark_dataset = tf.data.TFRecordDataset(LANDMARK_TFRECORDS_PATH_LIST)

image_feature_description = {
    'height': tf.io.FixedLenFeature([], tf.int64),
    'width': tf.io.FixedLenFeature([], tf.int64),
    'depth': tf.io.FixedLenFeature([], tf.int64),
    'info': tf.io.FixedLenFeature([17], tf.float32),
    'image_raw': tf.io.FixedLenFeature([], tf.string),
    }

def _read_tfrecord(serialized_example):

    example = tf.io.parse_single_example(serialized_example, image_feature_description)

    img = tf.image.decode_jpeg(example['image_raw'], channels = 3) # RGB rather than BGR!!! 
    img = (tf.cast(img, tf.float32) - 127.5) / 128.
    img_shape = [example['height'], example['width'], example['depth']]
    img = tf.reshape(img, img_shape)

    info = example['info']

    return img, info

parsed_pos_dataset = raw_pos_dataset.map(_read_tfrecord)
parsed_neg_dataset = raw_neg_dataset.map(_read_tfrecord)
parsed_part_dataset = raw_part_dataset.map(_read_tfrecord)
parsed_landmark_dataset = raw_landmark_dataset.map(_read_tfrecord)

parsed_image_dataset = tf.data.Dataset.zip((parsed_pos_dataset.repeat().shuffle(16384).batch(int(BATCH_SIZE * DATA_COMPOSE_RATIO[0])), 
                                            parsed_neg_dataset.repeat().shuffle(16384).batch(int(BATCH_SIZE * DATA_COMPOSE_RATIO[1])), 
                                            parsed_part_dataset.repeat().shuffle(16384).batch(int(BATCH_SIZE * DATA_COMPOSE_RATIO[2])), 
                                            parsed_landmark_dataset.repeat().shuffle(16384).batch(int(BATCH_SIZE * DATA_COMPOSE_RATIO[3]))))

def concatenate(pos_info, neg_info, part_info, landmark_info):

    img_tensor = tf.zeros((0, IMG_SIZE, IMG_SIZE, 3), dtype = tf.float32)
    label_tensor = tf.zeros((0, 17), dtype = tf.float32)
    pos_img = pos_info[0]
    neg_img = neg_info[0]
    part_img = part_info[0]
    landmark_img = landmark_info[0]
    pos_info = pos_info[1]
    neg_info = neg_info[1]
    part_info = part_info[1]
    landmark_info = landmark_info[1]
    img_tensor = tf.concat([img_tensor, pos_img, neg_img, part_img, landmark_img], axis = 0)
    info_tensor = tf.concat([label_tensor, pos_info, neg_info, part_info, landmark_info], axis = 0)

    return img_tensor, info_tensor

ds = parsed_image_dataset.map(concatenate)

Answer 1

您可以使用sample_from_datasets

根据指定的比例进行采样

有关示例，请参阅 this colab。我还复制了下面的代码。

import tensorflow as tf

batch_size = 8

pos = tf.data.Dataset.range(0, 100)
neg = tf.data.Dataset.range(100, 200)
part_face = tf.data.Dataset.range(200, 300)
landmark = tf.data.Dataset.range(300, 400)

dataset = tf.data.experimental.sample_from_datasets(
    [pos, neg, part_face, landmark], [1/7, 3/7, 1/7, 2/7])
dataset = dataset.batch(batch_size)

# Optionally shuffle data further. Samples will already be interspersed between datasets.
# dataset = dataset.map(lambda batch: tf.random.shuffle(batch))

for elem in dataset:
  print(elem.numpy())

# Prints
[200 300 100 201 301 302 101 303]
[  0 304 202 102 203 103 305 104]
[306 307 105 204 308 205 206   1]
[207 309 106 107 310 108 311 312]
[208 209 210   2 109 211 110 212]
...

TensorFlow 2-tf.keras：如何使用 tf.data API 和 TFRecords 训练像 MTCNN 这样的 tf.keras 多任务网络

TensorFlow 2-tf.keras: How to train a tf.keras multi-task network like MTCNN using tf.data API & TFRecords

multitasking

keras

tfrecord

tensorflow2.x