TensorFlow/Keras手动训练获取中间层输出并训练后续网络遇错求助

我是TensorFlow新手，尤其不熟悉手动训练流程，构建了三个自动编码器分别学习RGB图像的单个通道，现在想获取这些自动编码器的encoded层输出，拼接后传入另一个自动编码器训练。因GPU内存问题采用手动训练，但中间层输出是KerasTensor对象，无法获取数值用于批量训练，执行代码时报错。

自动编码器构建代码 #

def build_autoencoder():
    input_img = keras.Input(shape=(256, 256,1))
    x = Conv2D(32, kernel_size = (3,3), activation=LeakyReLU(alpha=0.03), padding = 'same')(input_img)
    x = MaxPooling2D((2, 2), padding='same')(x)
    x = Conv2D(16, kernel_size = (3,3), activation=LeakyReLU(alpha=0.03), padding = 'same')(x)
    x = Conv2D(8, kernel_size = (3,3), activation=LeakyReLU(alpha=0.03), padding = 'same')(x)
    encoded = MaxPooling2D((2, 2), padding='same', name='encoded')(x)
    x = Conv2D(8, (3, 3), activation=LeakyReLU(alpha=0.03), padding='same')(encoded)
    x = Conv2D(16, (3, 3), activation=LeakyReLU(alpha=0.03), padding='same')(x)
    x = UpSampling2D((2, 2))(x)
    x = Conv2D(32, (3, 3), activation=LeakyReLU(alpha=0.03), padding='same')(x)
    decoded = Conv2D(1, (3, 3), activation='sigmoid', padding='same')(x)

    autoencoder = keras.Model(inputs=input_img, outputs=decoded)
    autoencoder.compile(optimizer=Nadam(learning_rate=0.001), loss=tf.keras.losses.MeanSquaredError())
    return autoencoder

获取拼接编码输出的代码 #

encoded_red = model_red.get_layer('encoded').output
encoded_green = model_green.get_layer('encoded').output
encoded_blue = model_blue.get_layer('encoded').output
concatenated_encode = tf.keras.layers.Concatenate(axis=-1)([encoded_red, encoded_green, encoded_blue])
model_encoder_generator = build_encoder_generator(concatenated_encode)

手动训练代码 #

batch_size = 16
epochs=30

num_samples = len(y_train)
num_batches = num_samples // batch_size
optimizer = tf.keras.optimizers.Nadam(learning_rate=0.001)

batch_size = 16
epochs = 30

num_samples = len(y_train)
num_batches = num_samples // batch_size
concatenated_encode = tf.keras.layers.Concatenate(axis=-1)([encoded_red, encoded_green, encoded_blue])
optimizer = tf.keras.optimizers.Nadam(learning_rate=0.001)

history_encoder_generator = []
for epoch in range(epochs):
    epoch_loss = 0.0
    for batch_idx in range(0, num_samples, batch_size):
        batch_x = concatenated_encode[batch_idx : batch_idx + batch_size]
        batch_y = y_train[batch_idx : batch_idx + batch_size]

        with tf.GradientTape() as tape:
            predictions = model_encoder_generator(batch_x)
            loss = tf.keras.losses.mean_squared_error(batch_y, predictions)
        gradients = tape.gradient(loss, model_encoder_generator.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model_encoder_generator.trainable_variables))

        batch_mean_loss = tf.reduce_mean(loss)
        epoch_loss += batch_mean_loss
    avg_loss = epoch_loss / num_batches
    print("Epoch {}: Average Loss = {}".format(epoch + 1, avg_loss))
    history_encoder_generator.append(avg_loss)

报错信息 #

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[116], line 12
     10     predictions = model_encoder_generator(batch_x)
     11     loss = tf.keras.losses.mean_squared_error(batch_y, predictions)
---> 12 gradients = tape.gradient(loss, model_encoder_generator.trainable_variables)
     13 optimizer.apply_gradients(zip(gradients, model_encoder_generator.trainable_variables))
     15 batch_mean_loss = tf.reduce_mean(loss)

File ~\anaconda3\envs\test\lib\site-packages\tensorflow\python\eager\backprop.py:1113, in GradientTape.gradient(self, target, sources, output_gradients, unconnected_gradients)
   1107   output_gradients = (
   1108       composite_tensor_gradient.get_flat_tensors_for_gradients(
   1109           output_gradients))
   1110   output_gradients = [None if x is None else ops.convert_to_tensor(x)
   1111                       for x in output_gradients]
-> 1113 flat_grad = imperative_grad.imperative_grad(
   1114     self._tape,
   1115     flat_targets,
   1116     flat_sources,
   1117     output_gradients=output_gradients,
   1118     sources_raw=flat_sources_raw,
   1119     unconnected_gradients=unconnected_gradients)
   1121 if not self._persistent:
   1122   # Keep track of watched variables before setting tape to None
   1123   self._watched_variables = self._tape.watched_variables()

File ~\anaconda3\envs\test\lib\site-packages\tensorflow\python\eager\imperative_grad.py:67, in imperative_grad(tape, target, sources, output_gradients, sources_raw, unconnected_gradients)
     63 except ValueError:
     64   raise ValueError(
     65       "Unknown value for unconnected_gradients: %r" % unconnected_gradients)
---> 67 return pywrap_tfe.TFE_Py_TapeGradient(
     68     tape._tape,  # pylint: disable=protected-access
     69     target,
     70     sources,
     71     output_gradients,
     72     sources_raw,
     73     compat.as_str(unconnected_gradients.value))

AttributeError: 'KerasTensor' object has no attribute '_id'

核心问题分析 #

直接使用model.get_layer('encoded').output得到的是KerasTensor符号张量，它只是计算图的节点，不包含实际数据。手动训练时需要的是真实的批量数据，而不是符号张量，这才导致了报错。

正确步骤 #

1. 构建编码器模型（从预训练自动编码器中提取）

为每个预训练的自动编码器创建单独的编码器模型，用于生成实际的编码输出：

# 为三个自动编码器分别构建编码器模型
encoder_red = keras.Model(inputs=model_red.input, outputs=model_red.get_layer('encoded').output)
encoder_green = keras.Model(inputs=model_green.input, outputs=model_green.get_layer('encoded').output)
encoder_blue = keras.Model(inputs=model_blue.input, outputs=model_blue.get_layer('encoded').output)

2. 修改手动训练逻辑，实时生成编码输出

在批量训练时，对当前批次的RGB图像分通道输入到三个编码器，得到真实的编码张量后再拼接：

batch_size = 16
epochs = 30

num_samples = len(y_train)
num_batches = num_samples // batch_size
optimizer = tf.keras.optimizers.Nadam(learning_rate=0.001)

history_encoder_generator = []
for epoch in range(epochs):
    epoch_loss = 0.0
    # 假设x_train是原始RGB图像数据，形状为(样本数,256,256,3)
    for batch_idx in range(0, num_samples, batch_size):
        # 取出当前批次的RGB图像，并拆分通道
        batch_rgb = x_train[batch_idx : batch_idx + batch_size]
        batch_red = batch_rgb[..., 0:1]  # 提取红色通道，保持形状(16,256,256,1)
        batch_green = batch_rgb[..., 1:2]
        batch_blue = batch_rgb[..., 2:3]
        batch_y = y_train[batch_idx : batch_idx + batch_size]

        with tf.GradientTape() as tape:
            # 用三个编码器生成当前批次的编码输出
            encoded_red_batch = encoder_red(batch_red)
            encoded_green_batch = encoder_green(batch_green)
            encoded_blue_batch = encoder_blue(batch_blue)
            # 拼接编码输出
            concatenated_batch = tf.concat([encoded_red_batch, encoded_green_batch, encoded_blue_batch], axis=-1)
            # 传入后续模型得到预测
            predictions = model_encoder_generator(concatenated_batch)
            loss = tf.keras.losses.mean_squared_error(batch_y, predictions)
        
        gradients = tape.gradient(loss, model_encoder_generator.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model_encoder_generator.trainable_variables))

        batch_mean_loss = tf.reduce_mean(loss)
        epoch_loss += batch_mean_loss
    
    avg_loss = epoch_loss / num_batches
    print("Epoch {}: Average Loss = {}".format(epoch + 1, avg_loss))
    history_encoder_generator.append(avg_loss)

额外说明 #

• 如果不需要对三个预训练的自动编码器进行微调，可以在构建编码器模型时设置trainable=False，避免梯度回溯到这些模型，节省内存：

encoder_red.trainable = False
encoder_green.trainable = False
encoder_blue.trainable = False

• 确保model_encoder_generator的输入形状与拼接后的编码形状匹配，比如三个编码器的输出都是(64,64,8)，拼接后是(64,64,24)，则model_encoder_generator的输入应设为(64,64,24)。

内容的提问来源于stack exchange，提问作者WonderCoding