如何将非Flux生成的任意.BSON权重文件加载到Flux.jl中

你已经通过如下代码成功读取了BSON格式的权重字典：

weights = ONNX.load_weights("weights.bson")
# 读取到的权重为包含521个条目的字典，示例内容如下：
# Dict{String, Any} with 521 entries:
#   "Constant_123"  => fill(1.0)
#   "Constant_3300" => fill(2.0)
#   "Constant_2837" => fill(0)
#   ...

由于该权重是ONNX模型转换得到，并非Flux原生保存的模型文件，无法直接用Flux默认的加载接口导入，可按照以下步骤操作： #

步骤1：定义结构完全匹配的Flux模型 #

这是权重导入的前提，你需要先手动搭建Flux模型，保证模型的层顺序、每一层的参数数量、形状和原ONNX模型完全一致。

步骤2：建立权重映射关系并赋值 #

你读取到的权重字典key为ONNX节点命名，需要先和你定义的Flux模型参数建立一一对应关系，再完成赋值：

1. 先查看你定义的Flux模型的参数信息，确认每个参数的名称和形状：

# 遍历打印Flux模型所有参数的形状
for (idx, param) in enumerate(Flux.params(model))
    println("参数$idx 形状: $(size(param))")
end

2. 建立ONNX权重key到Flux参数的映射表，再批量赋值：

# 示例映射表，需根据你自己的模型对应关系替换内容
weight_mapping = Dict(
    "Constant_123" => model.conv1.weight,
    "Constant_3300" => model.conv1.bias,
    "Constant_2551" => model.bn1.ϵ,
    # 剩余500+参数依次补充映射关系
)

# 批量完成权重赋值
for (onnx_key, flux_param) in weight_mapping
    # 转换数据类型和Flux参数保持一致，默认Flux用Float32
    flux_param .= Float32.(weights[onnx_key])
end

步骤3：验证权重导入正确性 #

赋值完成后，取一组随机输入，分别传入原ONNX模型和导入权重后的Flux模型，核对两者输出的误差在1e-5以内即证明导入成功。 #

注意事项 #

• 若模型包含批量归一化、层归一化等带运行状态的层，除了权重和偏置外，还要把对应的移动均值、方差、eps等参数也赋值到Flux对应层的状态变量中
• 如果参数数量过多，可编写自动匹配脚本，按参数形状排序后自动关联对应，减少手动映射的工作量
• 若需要用GPU运行，赋值完成后调用model = model |> gpu即可把所有参数迁移到GPU设备

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