keras - TensorFlow 2-tf.keras:如何使用 tf.data API 和 TFRecords 训练 tf.keras 多任务网络,如 MTCNN
问题描述
最近我正在尝试使用 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)
解决方案
您可以使用sample_from_datasets根据指定的比率进行采样
有关示例,请参阅此 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]
...
推荐阅读
- json - 是否可以同时使用 Spring Boot 应用程序来控制网站和作为 API?
- windows - Windows Server Web Application Load Balancer - 如何检查它是使用粘性会话还是循环?
- c++ - 使用 void* 正确地将对象传递给 pthread_create()
- javascript - Javascript a = b 语句设置 b 值?
- r - 更改 bs4_book (bookdown) 中的右侧边栏标题?
- node.js - Feathers JS MongoDB 服务:动态集合名称
- amazon-web-services - 多设备和时间戳的最佳分区方法
- angular - 如何从Angular中的多选下拉菜单中汇总多个值
- swift - 关联的枚举表示为自定义字体的字符串和 CGFloat?
- reactjs - 在 React 中使用数据 URI 而不是小于 10kb 的图像的路径的原因是什么?