重训练模型时加载权重报错：axes don't match array 求助

问题重现 #

执行权重加载代码时触发报错：

model.load_weights(config.training.weights_path, by_name=True) 

报错栈：

File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 2232, in load_weights
    f, self.layers, skip_mismatch=skip_mismatch)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/hdf5_format.py", line 761, in load_weights_from_hdf5_group_by_name
    layer, weight_values, original_keras_version, original_backend)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/hdf5_format.py", line 405, in preprocess_weights_for_loading
    weights[0] = np.transpose(weights[0], (3, 2, 0, 1))
File "<__array_function__ internals>", line 6, in transpose
File "/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py", line 653, in transpose
    return _wrapfunc(a, 'transpose', axes)
File "/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py", line 58, in _wrapfunc
    return bound(*args, **kwds)
ValueError: axes don't match array

已尝试将模型文件后缀从.h5改为.hdf5，无效。

报错原因 #

本质是待加载的权重张量维度与当前模型对应层的张量维度不匹配，Keras在尝试按规则转置权重时，发现轴的数量/大小和目标张量不匹配。常见诱因：

• 当前模型的层结构（如卷积层过滤器数量、输入通道数、全连接层单元数）和导出权重的原始模型不一致
• 权重文件由不同框架/不同版本的Keras导出（如原生Keras vs TF-Keras），存在通道顺序（channels_first/channels_last）差异
• 启用by_name=True按层名匹配，但对应层的参数形状已修改

解决方案 #

1. 严格对齐模型结构

确保当前定义的模型和导出权重的原始模型完全一致，包括每一层的参数细节。可以打印层信息对比：

# 遍历打印当前模型的层名称和输出形状
for layer in model.layers:
    print(f"Layer: {layer.name}, Output Shape: {layer.output_shape}")

如果是做迁移学习需要修改模型结构，不要直接用load_weights(by_name=True)，建议先加载完整原始模型，再替换顶层结构，或手动加载匹配层的权重。

2. 修复框架/版本的维度差异

如果权重来自不同框架或Keras版本，先统一通道顺序：

from tensorflow.keras import backend as K
# 设置和原始模型一致的通道顺序，比如channels_last或channels_first
K.set_image_data_format('channels_last')

若仍无效，可手动读取权重并调整维度后赋值：

import h5py
import numpy as np

with h5py.File(config.training.weights_path, 'r') as weight_file:
    # 示例：读取conv1层的卷积核权重
    conv_kernel = weight_file['conv1']['kernel:0'][()]
    # 根据当前模型层的维度要求调整轴顺序，示例为(3,3,3,64)转(64,3,3,3)
    adjusted_kernel = np.transpose(conv_kernel, (3, 2, 0, 1))
    # 获取对应层并设置权重（需包含偏置等其他参数）
    model.get_layer('conv1').set_weights([adjusted_kernel, weight_file['conv1']['bias:0'][()]])

3. 跳过不匹配的层（谨慎使用）

如果确认部分层不需要加载权重，可添加skip_mismatch=True参数跳过维度不匹配的层：

model.load_weights(config.training.weights_path, by_name=True, skip_mismatch=True)

注意：该参数会忽略不匹配层的权重加载，可能导致模型初始化不完整，仅适用于明确不需要这些层权重的场景。

4. 重新导出权重文件

若能获取权重对应的原始模型，使用当前环境的TensorFlow/Keras版本重新导出权重，彻底避免版本兼容问题。

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