Оценщик тензора: узкие места в кэше

Следуя руководству по классификации изображений в тензорном потоке, сначала он кэширует узкое место каждого изображения:

def: cache_bottlenecks())

Я переписал тренировку с помощью тензометрического Estimator. Это действительно упростило весь код. Однако я хочу кэшировать функции узких мест здесь.

Вот моя model_fn. Я хочу, чтобы кешировать результаты dense слоя, чтобы я мог вносить изменения в фактическое обучение без необходимости вычислять узкие места каждый раз.

Как я могу это сделать?

def model_fn(features, labels, mode, params):
    is_training = mode == tf.estimator.ModeKeys.TRAIN

    num_classes = len(params['label_vocab'])

    module = hub.Module(params['module_spec'], trainable=is_training and params['train_module'])
    bottleneck_tensor = module(features['image'])

    with tf.name_scope('final_retrain_ops'):
        logits = tf.layers.dense(bottleneck_tensor, units=num_classes, trainable=is_training)  # save this?

    def train_op_fn(loss):
        optimizer = tf.train.AdamOptimizer()
        return optimizer.minimize(loss, global_step=tf.train.get_global_step())

    head = tf.contrib.estimator.multi_class_head(n_classes=num_classes, label_vocabulary=params['label_vocab'])

    return head.create_estimator_spec(
        features, mode, logits, labels, train_op_fn=train_op_fn
    )

Ответ 1

TF не может работать с кодом. Вам следует:

Экспортное узкое место для файла из сырой сети.
Используйте результат узкого места в качестве входных данных, используйте другую сеть для обучения своих данных.

Ответ 2

Подробнее о том, что сказал @Feng:

см. TFRecords и TFExamples и Загрузить изображения

Примерно так должно работать (не проверено):

# Serialize the data into two tfrecord files
tf.enable_eager_execution()
feature_extractor = ...
features_file = tf.python_io.TFRecordWriter('features.tfrec')
label_file = tf.python_io.TFRecordWriter('labels.tfrec')

for images, labels in dataset:
  features = feature_extractor(images)
  features_file.write(tf.serialize_tensor(features))
  label_file.write(tf.serialize_tensor(labels))

# Parse the files and zip them together
def parse(type, shape):
  _def parse(x):
    result = tf.parse_tensor(x, out_type=shape)
    result = tf.reshape(result, FEATURE_SHAPE)
    return result
  return parse

features_ds = tf.data.TFRecordDataset('features.tfrec')
features_ds = features_ds.map(parse(tf.float32, FEATURE_SHAPE), num_parallel_calls=AUTOTUNE)

labels_ds = tf.data.TFRecordDataset('labels.tfrec')
labels_ds = labels_ds.map(parse(tf.float32, FEATURE_SHAPE), num_parallel_calls=AUTOTUNE)

ds = tf.data.Dataset.zip(features_ds, labels_ds)
ds = ds.unbatch().shuffle().repeat().batch().prefetch()...

Вы также можете сделать это, используя Dataset.cache, но я не уверен на 100% в деталях.