如何切割预训练的TensorFlow模型？附磁盘存储自编码器模型示例

要切割你这个预训练的自编码器，把它拆成独立的编码器和解码器部分，核心就是提取原模型中对应子结构的权重和计算图，下面分两种常见场景给你具体步骤：

一、如果模型是SavedModel格式（TensorFlow默认推荐格式） #

假设你已经把完整的自编码器保存成了SavedModel，先加载整个模型：

import tensorflow as tf
lays = tf.keras.layers

# 加载预训练的完整自编码器
autoencoder_model = tf.keras.models.load_model("path/to/your/autoencoder")

1. 提取编码器 #

你的编码器是前三个conv2d层，对应原模型中到net = lays.conv2d(net, 8, [5,5], stride=4, padding='SAME')这一步的输出。你可以通过两种方式定义子模型：

# 方式1：通过层索引（注意：原模型layers列表里，Input层是第0个，三个conv2d分别是第1、2、3个）
encoder = tf.keras.Model(
    inputs=autoencoder_model.input,
    outputs=autoencoder_model.layers[3].output
)

# 方式2：通过层名称（更可靠，避免结构变动导致索引失效）
# 可以先打印autoencoder_model.layers查看各层名称，假设第三个conv2d的名称是"conv2d_2"
encoder = tf.keras.Model(
    inputs=autoencoder_model.input,
    outputs=autoencoder_model.get_layer("conv2d_2").output
)

2. 提取解码器 #

解码器是从第一个conv2d_transpose层开始的部分，需要基于编码器的输出形状创建新输入，再串联原模型的后续层：

# 获取编码器输出的形状（去掉batch维度）
encoder_out_shape = encoder.output_shape[1:]
# 创建解码器的输入层
decoder_input = tf.keras.Input(shape=encoder_out_shape)

# 串联原模型中解码器的所有层
x = decoder_input
# 原模型中从第4层开始是解码器的层（前3个conv2d是编码器，第0个是Input层）
for layer in autoencoder_model.layers[4:]:
    x = layer(x)

# 定义解码器模型
decoder = tf.keras.Model(inputs=decoder_input, outputs=x)

3. 验证切割是否正确 #

用一张测试图片验证输出是否一致：

# 替换成你的实际输入形状，比如(1, 64, 64, 3)
test_img = tf.random.normal((1, 64, 64, 3))  

# 完整模型输出
full_output = autoencoder_model(test_img)
# 切割后编码器+解码器的输出
encoded = encoder(test_img)
decoded = decoder(encoded)

# 检查输出是否在浮点误差范围内一致
print(tf.reduce_mean(tf.abs(full_output - decoded)))  # 结果应该接近0

二、如果模型仅保存了权重文件（.ckpt格式） #

如果只有权重文件，需要先重新定义编码器和解码器的结构，再把原模型的权重复制过去：

1. 重新定义编码器和解码器结构 #

# 定义编码器
def build_encoder(input_shape):
    inputs = tf.keras.Input(shape=input_shape)
    net = lays.conv2d(inputs, 32, [5,5], stride=2, padding='SAME')
    net = lays.conv2d(net, 16, [5,5], stride=2, padding='SAME')
    net = lays.conv2d(net, 8, [5,5], stride=4, padding='SAME')
    return tf.keras.Model(inputs=inputs, outputs=net)

# 定义解码器（补全你原模型中未写完的conv2d_transpose层，示例假设最后输出3通道图像）
def build_decoder(encoder_out_shape):
    inputs = tf.keras.Input(shape=encoder_out_shape)
    net = lays.conv2d_transpose(inputs, 16, [5,5], stride=4, padding='SAME')
    net = lays.conv2d_transpose(net, 32, [5,5], stride=2, padding='SAME')
    net = lays.conv2d_transpose(net, 3, [5,5], stride=2, padding='SAME', activation='sigmoid')
    return tf.keras.Model(inputs=inputs, outputs=net)

2. 加载权重到子模型 #

# 先加载完整模型（如果有结构文件），或者先定义完整模型再加载权重
autoencoder_model = tf.keras.models.load_model("path/to/autoencoder")
# 如果只有权重文件：
# autoencoder_model = autoencoder(tf.keras.Input(shape=(64,64,3)))
# autoencoder_model.load_weights("path/to/weights.ckpt")

# 构建编码器并复制权重
encoder = build_encoder(input_shape=(64,64,3))  # 替换成你的实际输入形状
# 原模型第1-3层对应编码器的第1-3层
for i in range(1, 4):
    encoder.layers[i].set_weights(autoencoder_model.layers[i].get_weights())

# 构建解码器并复制权重
encoder_out_shape = encoder.output_shape[1:]
decoder = build_decoder(encoder_out_shape)
# 原模型第4层及以后对应解码器的第1层及以后
for i in range(1, len(decoder.layers)):
    decoder.layers[i].set_weights(autoencoder_model.layers[i+3].get_weights())

一些注意事项 #

• 优先用层名称：尽量通过层名称提取层，避免因模型结构微调导致索引失效。
• 自定义层处理：如果原模型用了自定义层，加载时要确保用tf.keras.utils.get_custom_objects()注册，否则load_model会报错。
• 保存子模型：切割完成后，用encoder.save("encoder_savedmodel")和decoder.save("decoder_savedmodel")保存成SavedModel，方便后续直接使用。

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