Производная второго порядка Tensorflow 2 после примененных градиентов в стремительном исполнении

Question

Я пытаюсь реализовать модель Agnosti c Meta Learning (MAML) в Tensorflow 2. Для этого алгоритма вычисляется производная второго порядка в виде:

for multiple sets, each with an individual model:
  1. Determine the loss for a training set
  2. Determine the gradients for this loss w.r.t. the model
  3. Apply the gradients (theta' <- theta-alpha*gradients)

4. Determine the loss for an evaluation set for each of the updated models
5. Determine the gradient of the sum of the losses (4) w.r.t. the global model

Я пытаюсь реализовать это в тензорном потоке 2 с GradientTape, но вторым градиентом всегда является массив None с.

Моя реализация выглядит так:

with tf.GradientTape() as meta_update_tape:
    # inner loop (for all tasks)
    for bi in range(batch_size):

       # the training and evaluation batches
       x_i, y_i, x_i_prime, y_i_prime = dataset.batch_with_eval()

       # reset the weights to to current global weights before training for this batch
       model_copy = copy_model_weights(source=model, target=model_copy)

       # Compute loss using theta_global for D_i
       with tf.GradientTape() as inner_update_tape:
           inner_loss, _, _ = compute_loss(model_copy, x_i, y_i)

       gradients_inner_update = inner_update_tape.gradient(inner_loss, model_copy.trainable_variables)

       # update model parameters (apply theta_i_prime)
       # inner_optimizer.apply_gradients(zip(gradients_inner_update, model_copy.trainable_variables))
       conv_layers  = [i for i,var in enumerate(model_copy.trainable_variables) if var.name.startswith('conv2d')]
       for layer in conv_layers:
           model_copy.trainable_variables[layer].assign(tf.subtract(model.trainable_variables[layer], tf.multiply(0.4, gradients_inner_update[layer])))

       # calculate loss with theta_i_prime with eval set
       loss_eval, accuracy_eval, _ = compute_loss(model_copy, x_i_prime, y_i_prime)

       batch_losses.append(loss_eval)

   sum_losses = tf.reduce_sum(batch_losses) / tf.cast(batch_size, dtype=tf.float32)

# calculate gradient over all losses w.r.t global theta (trainable variables for global model)
gradients_meta_update = meta_update_tape.gradient(sum_losses, model.trainable_variables)

# gradients_meta_update is [None, None, None, ...]

Я пробовал:

• для запуска без l oop (только одна партия)
• для получения суммы потерь вне внешней ленты
• для применения градиентов для внутренней обновление с оптимизатором SGD (это то, что я бы предпочел сделать)
• чтобы вообще не применять внутренний градиент
• иметь внутреннюю ленту с опцией persistent=True
• явно смотрите model.trainable_variables на внешней ленте

У меня заканчиваются идеи, и документация мне не помогает, поэтому я ценю каждое предложение. Спасибо!