如何修改Keras预训练模型输入维度且不丢失已训练权重

核心结论 #

不可以直接仅修改输入层维度就完成切换，不过可以通过两种方案实现需求，全程不需要调整输入图像，也不会损失已训练好的权重。你要修改的输入维度对应图像数据格式从TensorFlow默认的通道最后(channels_last, [H,W,C])切换为通道优先(channels_first, [C,H,W])，适配两种场景的方案如下：

方案1：插入维度置换层（最简单，无需修改原有层参数/权重） #

该方案完全保留原有模型的所有层配置和权重，仅在输入层后增加一个维度转换层，自动把你输入的(None,3,224,224)格式转成原有模型需要的格式，代码如下：

import keras

# 加载原有模型
model = keras.models.load_model('/content/Sample_MobileNetV2_7Class_210721.hdf5', compile=False)

# 新建符合你要求的输入层
new_input = keras.layers.Input(shape=(3, 224, 224))
# 新增维度置换层，将[C,H,W]转成原有模型需要的[H,W,C]
permute_layer = keras.layers.Permute((2,3,1))(new_input)
# 接入原有模型的所有层
output = model(permute_layer)

# 构建新模型
new_model = keras.Model(inputs=new_input, outputs=output)
# 验证维度
new_model.summary()

验证时你输入(None,3,224,224)格式的图像即可正常推理，原有权重100%保留。

方案2：完全切换模型为通道优先格式 #

如果需要整个模型都适配通道优先的计算逻辑，不需要插入转换层，可以按以下步骤操作：

1. 先设置后端数据格式

import keras
from keras import backend as K
K.set_image_data_format('channels_first')

2. 重构模型结构，遍历原有层修改配置，调整权重维度

# 加载原有模型
old_model = keras.models.load_model('/content/Sample_MobileNetV2_7Class_210721.hdf5', compile=False)

# 新建输入层
new_input = keras.layers.Input(shape=(3,224,224))
x = new_input

for layer in old_model.layers[1:]: # 跳过原有输入层
    # 处理涉及空间维度的层，修改data_format参数
    if isinstance(layer, (keras.layers.Conv2D, keras.layers.BatchNormalization, keras.layers.MaxPooling2D, keras.layers.AveragePooling2D, keras.layers.ReLU)):
        # 复制原有层配置，修改data_format
        config = layer.get_config()
        config['data_format'] = 'channels_first'
        new_layer = type(layer).from_config(config)
        x = new_layer(x)
        # 调整卷积层权重维度适配通道优先
        if isinstance(layer, keras.layers.Conv2D):
            # 原有权重维度[kernel_h, kernel_w, in_channels, out_channels] -> 通道优先需要[out_channels, in_channels, kernel_h, kernel_w]
            old_weights = layer.get_weights()
            new_kernel = old_weights[0].transpose(3,2,0,1)
            new_layer.set_weights([new_kernel, old_weights[1]] if len(old_weights) >1 else [new_kernel])
        else:
            # 其他层权重直接复用
            new_layer.set_weights(layer.get_weights())
    else:
        # 不涉及空间维度的层直接复用
        x = layer(x)

# 构建新模型
new_model = keras.Model(inputs=new_input, outputs=x)
# 验证维度
new_model.summary()

注意事项 #

• 方案1仅新增了一个无参数的Permute层，计算开销可以忽略，优先推荐使用
• 方案2需要确保所有涉及空间维度的层都修改了data_format参数，否则会出现维度不匹配错误

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