最近我正在尝试使用 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。
因此,例如,在我制作了四个 TFRecord 样本后,我需要从 postfrecords 中提取 128 个样本,从 neg 中提取 384 个样本,从 part-face 中提取 128 个样本,从界标 tfrecords 中提取 256 个样本,并将它们混合成一个 mini-batch . 然后我需要在训练前对批次进行洗牌。
在使用 TFRecords 和 tf.data API 时,我真的不知道如何进行此操作。现在我只能通过批量读取图像和标签来实现所有这些步骤,但是训练速度太慢了。所以我想知道是否有任何有效的方法可以做到这一点。
任何建议表示赞赏。
更新于 2020.04.04 15:38
感谢@AAudibert,我认为他/她的回答效果很好,我也想出了一种实现方法。这是供参考的代码:
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)