假设我们正在构建一个基本的 CNN 来识别猫和狗的图片（二元分类器）。
此类 CNN 的示例如下：
model = Sequential([
Conv2D(32, (3,3), input_shape=...),
Activation('relu'),
MaxPooling2D(pool_size=(2,2),
Conv2D(32, (3,3), input_shape=...),
Activation('relu'),
MaxPooling2D(pool_size=(2,2)
Conv2D(64, (3,3), input_shape=...),
Activation('relu'),
MaxPooling2D(pool_size=(2,2),
Flatten(),
Dense(64),
Activation('relu'),
Dropout(0.5),
Dense(1),
Activation('sigmoid')
])
我们还假设我们希望将模型分成两部分，或两个模型，称为model_0 and model_1.
model_0将处理输入，并且model_1将采取model_0输出并将其作为输入。
例如，之前的模型将变为：
model_0 = Sequential([
Conv2D(32, (3,3), input_shape=...),
Activation('relu'),
MaxPooling2D(pool_size=(2,2),
Conv2D(32, (3,3), input_shape=...),
Activation('relu'),
MaxPooling2D(pool_size=(2,2)
Conv2D(64, (3,3), input_shape=...),
Activation('relu'),
MaxPooling2D(pool_size=(2,2)
])
model_1 = Sequential([
Flatten(),
Dense(64),
Activation('relu'),
Dropout(0.5),
Dense(1),
Activation('sigmoid')
])
如何将这两个模型当作一个模型来训练？我尝试手动设置渐变，但我不明白如何传递渐变model_1 to model_0:
for epoch in range(epochs):
for step, (x_batch, y_batch) in enumerate(train_generator):
# model 0
with tf.GradientTape() as tape_0:
y_pred_0 = model_0(x_batch, training=True)
# model 1
with tf.GradientTape() as tape_1:
y_pred_1 = model_1(y_pred_0, training=True)
loss_value = loss_fn(y_batch_tensor, y_pred_1)
grads_1 = tape_1.gradient(y_pred_1, model_1.trainable_weights)
grads_0 = tape_0.gradient(y_pred_0, model_0.trainable_weights)
optimizer.apply_gradients(zip(grads_1, model_1.trainable_weights))
optimizer.apply_gradients(zip(grads_0, model_0.trainable_weights))
这种方法当然行不通，因为我基本上只是分别训练两个模型并将它们绑定起来，这不是我想要实现的。
这是一个针对此问题的简单版本的 Google Colab 笔记本，仅使用两个全连接层和两个激活函数：https://colab.research.google.com/drive/14Px1rJtiupnB6NwtvbgeVYw56N1xM6JU#scrollTo=PeqtJJWS3wyG https://colab.research.google.com/drive/14Px1rJtiupnB6NwtvbgeVYw56N1xM6JU#scrollTo=PeqtJJWS3wyG
请注意，我知道Sequential([model_0, model_1])，但这不是我想要达到的目标。我想手动执行反向传播步骤。
另外，我想继续使用两个单独的磁带。这里的技巧是使用grads_1计算grads_0.
有什么线索吗？
在寻求帮助并更好地了解动态之后自动微分 (or autodiff），我设法得到了一个可行的、简单的例子来说明我想要实现的目标。尽管这种方法不能完全解决问题，但它使我们在理解如何解决当前问题方面向前迈进了一步。
我已将模型简化为更小的模型：
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Activation, Dense, Layer, Flatten, Conv2D
import numpy as np
tf.random.set_seed(0)
# 3 batches, 10x10 images, 1 channel
x = tf.random.uniform((3, 10, 10, 1))
y = tf.cast(tf.random.uniform((3, 1)) > 0.5, tf.float32)
layer_0 = Sequential([Conv2D(filters=6, kernel_size=2, activation="relu")])
layer_1 = Sequential([Conv2D(filters=6, kernel_size=2, activation="relu")])
layer_2 = Sequential([Flatten(), Dense(1), Activation("sigmoid")])
loss_fn = tf.keras.losses.MeanSquaredError()
我们将其分为三个部分，layer_0, layer_1, layer_2。普通方法只是将所有内容放在一起并一一计算梯度（或一步计算）：
with tf.GradientTape(persistent=True) as tape:
out_layer_0 = layer_0(x)
out_layer_1 = layer_1(out_layer_0)
out_layer_2 = layer_2(out_layer_1)
loss = loss_fn(y, out_layer_2)
只需简单调用即可计算不同的梯度tape.gradient:
ref_conv_dLoss_dWeights2 = tape.gradient(loss, layer_2.trainable_weights)
ref_conv_dLoss_dWeights1 = tape.gradient(loss, layer_1.trainable_weights)
ref_conv_dLoss_dWeights0 = tape.gradient(loss, layer_0.trainable_weights)
ref_conv_dLoss_dY = tape.gradient(loss, out_layer_2)
ref_conv_dLoss_dOut1 = tape.gradient(loss, out_layer_1)
ref_conv_dOut2_dOut1 = tape.gradient(out_layer_2, out_layer_1)
ref_conv_dLoss_dOut0 = tape.gradient(loss, out_layer_0)
ref_conv_dOut1_dOut0 = tape.gradient(out_layer_1, out_layer_0)
ref_conv_dOut0_dWeights0 = tape.gradient(out_layer_0, layer_0.trainable_weights)
ref_conv_dOut1_dWeights1 = tape.gradient(out_layer_1, layer_1.trainable_weights)
ref_conv_dOut2_dWeights2 = tape.gradient(out_layer_2, layer_2.trainable_weights)
稍后我们将使用这些值来比较我们方法的正确性。
对于分裂，我们的意思是每个layer_x需要有自己的GradientTape，负责生成自己的梯度：
with tf.GradientTape(persistent=True) as tape_0:
out_layer_0 = model.layers[0](x)
with tf.GradientTape(persistent=True) as tape_1:
tape_1.watch(out_layer_0)
out_layer_1 = model.layers[1](out_layer_0)
with tf.GradientTape(persistent=True) as tape_2:
tape_2.watch(out_layer_1)
out_flatten = model.layers[2](out_layer_1)
out_layer_2 = model.layers[3](out_flatten)
loss = loss_fn(y, out_layer_2)
现在，简单地使用tape_n.gradient因为每一步都行不通。我们基本上丢失了很多事后无法恢复的信息。
Instead, we have to use tape.jacobian https://www.tensorflow.org/api_docs/python/tf/GradientTape#jacobian and tape.batch_jacobian https://www.tensorflow.org/api_docs/python/tf/GradientTape#batch_jacobian, except for , as we only have one value as a source.
dOut0_dWeights0 = tape_0.jacobian(out_layer_0, model.layers[0].trainable_weights)
dOut1_dOut0 = tape_1.batch_jacobian(out_layer_1, out_layer_0)
dOut1_dWeights1 = tape_1.jacobian(out_layer_1, model.layers[1].trainable_weights)
dOut2_dOut1 = tape_2.batch_jacobian(out_layer_2, out_layer_1)
dOut2_dWeights2 = tape_2.jacobian(out_layer_2, model.layers[3].trainable_weights)
dLoss_dOut2 = tape_2.gradient(loss, out_layer_2) # or dL/dY
我们将使用几个实用函数来将结果调整为我们想要的：
def add_missing_axes(source_tensor, target_tensor):
len_missing_axes = len(target_tensor.shape) - len(source_tensor.shape)
# note: the number of tf.newaxis is determined by the number of axis missing to reach
# the same dimension of the target tensor
assert len_missing_axes >= 0
# convenience renaming
source_tensor_extended = source_tensor
# add every missing axis
for _ in range(len_missing_axes):
source_tensor_extended = source_tensor_extended[..., tf.newaxis]
return source_tensor_extended
def upstream_gradient_loss_weights(dOutUpstream_dWeightsLocal, dLoss_dOutUpstream):
dLoss_dOutUpstream_extended = add_missing_axes(dLoss_dOutUpstream, dOutUpstream_dWeightsLocal)
# reduce over the first axes
len_reduce = range(len(dLoss_dOutUpstream.shape))
return tf.reduce_sum(dOutUpstream_dWeightsLocal * dLoss_dOutUpstream_extended, axis=len_reduce)
def upstream_gradient_loss_out(dOutUpstream_dOutLocal, dLoss_dOutUpstream):
dLoss_dOutUpstream_extended = add_missing_axes(dLoss_dOutUpstream, dOutUpstream_dOutLocal)
len_reduce = range(len(dLoss_dOutUpstream.shape))[1:]
return tf.reduce_sum(dOutUpstream_dOutLocal * dLoss_dOutUpstream_extended, axis=len_reduce)
最后，我们可以应用链式法则：
dOut2_dOut1 = tape_2.batch_jacobian(out_layer_2, out_layer_1)
dOut2_dWeights2 = tape_2.jacobian(out_layer_2, model.layers[3].trainable_weights)
dLoss_dOut2 = tape_2.gradient(loss, out_layer_2) # or dL/dY
dLoss_dWeights2 = upstream_gradient_loss_weights(dOut2_dWeights2[0], dLoss_dOut2)
dLoss_dBias2 = upstream_gradient_loss_weights(dOut2_dWeights2[1], dLoss_dOut2)
dLoss_dOut1 = upstream_gradient_loss_out(dOut2_dOut1, dLoss_dOut2)
dLoss_dWeights1 = upstream_gradient_loss_weights(dOut1_dWeights1[0], dLoss_dOut1)
dLoss_dBias1 = upstream_gradient_loss_weights(dOut1_dWeights1[1], dLoss_dOut1)
dLoss_dOut0 = upstream_gradient_loss_out(dOut1_dOut0, dLoss_dOut1)
dLoss_dWeights0 = upstream_gradient_loss_weights(dOut0_dWeights0[0], dLoss_dOut0)
dLoss_dBias0 = upstream_gradient_loss_weights(dOut0_dWeights0[1], dLoss_dOut0)
print("dLoss_dWeights2 valid:", tf.experimental.numpy.allclose(ref_conv_dLoss_dWeights2[0], dLoss_dWeights2).numpy())
print("dLoss_dBias2 valid:", tf.experimental.numpy.allclose(ref_conv_dLoss_dWeights2[1], dLoss_dBias2).numpy())
print("dLoss_dWeights1 valid:", tf.experimental.numpy.allclose(ref_conv_dLoss_dWeights1[0], dLoss_dWeights1).numpy())
print("dLoss_dBias1 valid:", tf.experimental.numpy.allclose(ref_conv_dLoss_dWeights1[1], dLoss_dBias1).numpy())
print("dLoss_dWeights0 valid:", tf.experimental.numpy.allclose(ref_conv_dLoss_dWeights0[0], dLoss_dWeights0).numpy())
print("dLoss_dBias0 valid:", tf.experimental.numpy.allclose(ref_conv_dLoss_dWeights0[1], dLoss_dBias0).numpy())
输出将是：
dLoss_dWeights2 valid: True
dLoss_dBias2 valid: True
dLoss_dWeights1 valid: True
dLoss_dBias1 valid: True
dLoss_dWeights0 valid: True
dLoss_dBias0 valid: True
因为所有值都彼此接近。请注意，使用雅可比行列式的方法，我们将有一定程度的误差/近似值，大约10^-7，但我认为这已经足够好了。
对于极限模型或玩具模型来说，这是完美的并且效果很好。然而，在实际场景中，您将拥有大量尺寸的大图像。在处理雅可比行列式时，这并不理想，因为雅可比行列式可以很快达到很高的维度。但这本身就是一个问题。
您可以在以下资源中阅读有关该主题的更多信息：