TensorFlow: Какой самый простой способ включить прогнозы из одной модели в обучение новой модели?

Каков самый простой способ использовать tf.estimator подготовленную модель A во время обучения другой модели B?

Веса в модели A фиксированы. В модели B я хотел бы взять некоторые входы, вычислить, передать эти результаты в модель A, а затем сделать еще несколько вычислений на выходе.

Простой пример:

ModelA возвращает tf.matmul(input,weights)

В ModelB я хотел бы сделать следующее:

  x1 = tf.matmul(new_inputs,new_weights1)
  x2 = modelA(x1) # with fixed weights
  return tf.matmul(x2,new_weights2)

Но с более сложными моделями A и B, каждый из которых обучается как tf.estimator (хотя я рад не использовать оценки, если есть другое простое решение - я использую их, потому что я хотел бы использовать ML Engine),

Этот вопрос связан, но предлагаемое решение не работает для обучения модели B, поскольку градиенты tf.py_func равны [None]. Я попытался зарегистрировать градиент для tf.py_func, но это не с

Неподдерживаемый тип объекта Тензор

Я также попробовал tf.import_graph_def для модели A, но это, похоже, загружает предварительно обработанный график, но не фактические веса.