如何将PyTorch中拆分3层的Conv2D实现正确迁移至Keras/TensorFlow

原Keras实现的核心错误 #

• 变量名直接覆盖：实例化完Keras层对象后，直接将first_layer、last_layer变量赋值为普通张量，完全丢失了之前定义的层对象。后续拼接的new_layers列表中全是基础张量，根本不是可参与计算的Keras层，自然无法被模型识别、存储权重。
• 层类型误用：将PyTorch中的普通1x1 Conv2D替换成了SeparableConv2D，这类层的内部结构是“逐通道卷积+1x1逐点卷积”，和需要实现的独立单层卷积结构完全不匹配，权重形状也无法对齐。
• 参数非法：给core_layer的padding参数传了元组(1,1)，Keras的Conv2D的padding参数仅接受"valid"或"same"两个字符串选项，传数值元组会直接触发参数校验错误。
• 初始化逻辑缺失：Keras层是延迟初始化机制，只有接入计算图、完成build流程后才会生成对应的权重变量。原代码没有构建计算图触发层初始化，提前准备的权重张量根本没有赋值入口。
• 维度顺序不匹配：PyTorch卷积层的权重格式为[输出通道, 输入通道, 卷积核高, 卷积核宽]，Keras默认的权重格式为[卷积核高, 卷积核宽, 输入通道, 输出通道]，直接做简单转置、升维不调整维度顺序，就算能赋值，计算结果也会完全错误。

正确的迁移实现步骤 #

• 统一用普通Conv2D定义三层结构，参数和PyTorch完全对齐：
  • 首层：3输出通道、1x1卷积、stride=1、padding=valid、无偏置
  • 中间层：16输出通道、3x3卷积、stride=1、padding=same、无偏置
  • 末层：64输出通道、1x1卷积、stride=1、padding=valid、带偏置
• 先通过输入层串联三层结构，触发Keras的层初始化，完成计算图构建
• 按照Keras的权重维度格式，转换准备好的原始权重张量：
  • 首层3x3的原始矩阵，转置后升为[1, 1, 输入通道=3, 输出通道=3]的形状
  • 中间层形状为[16, 3, 3, 3]的原始张量，调整维度顺序为[3, 3, 输入通道=3, 输出通道=16]
  • 末层形状为(64,16)的原始矩阵，升维为[1, 1, 输入通道=16, 输出通道=64]的形状，偏置直接复用原PyTorch层的偏置张量
• 用层自带的set_weights方法完成权重赋值，不要直接修改层内部的变量。

可直接运行的参考代码 #

import tensorflow as tf

# 定义与PyTorch完全对齐的三层卷积结构
first_layer = tf.keras.layers.Conv2D(
    filters=3, kernel_size=1, strides=1,
    padding="valid", dilation_rate=1, use_bias=False
)
core_layer = tf.keras.layers.Conv2D(
    filters=16, kernel_size=3, strides=1,
    padding="same", dilation_rate=1, use_bias=False
)
last_layer = tf.keras.layers.Conv2D(
    filters=64, kernel_size=1, strides=1,
    padding="valid", dilation_rate=1, use_bias=True
)

# 构建计算图，触发层初始化
input_tensor = tf.keras.Input(shape=(None, None, 3))
x = first_layer(input_tensor)
x = core_layer(x)
x = last_layer(x)
model = tf.keras.Model(inputs=input_tensor, outputs=x)

# 权重维度转换（对齐PyTorch与Keras的格式差异）
# 假设已加载好原始张量：first(3x3)、core([16,3,3,3])、last(64x16)、original_layer_bias
first_weight = tf.transpose(first, perm=[1, 0])[tf.newaxis, tf.newaxis, ...]
core_weight = tf.transpose(core, perm=[2, 3, 1, 0])
last_weight = last[tf.newaxis, tf.newaxis, ...]

# 赋值权重
first_layer.set_weights([first_weight])
core_layer.set_weights([core_weight])
last_layer.set_weights([last_weight, original_layer_bias])

# 验证：此时model.weights可正常返回所有层的权重，前向计算结果与PyTorch实现完全对齐

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