如何在Flux.jl中加载使用.pth格式的PyTorch模型？

首先给出明确结论：无法直接加载 .pth 格式的PyTorch模型到Flux.jl中使用。主要原因是：

• .pth是基于Python pickle协议的序列化文件，存储的是PyTorch专属的张量、算子对象，和Flux的Julia原生层结构完全不兼容
• Flux官方目前没有提供直接解析.pth格式的工具

以下是两种可落地的替代方案： #

方案1：ONNX格式中转（适用于模型结构均为标准算子的场景） #

这是成本最低的迁移方式，借助通用的开放神经网络交换格式ONNX做中间转换：

1. 先在PyTorch侧导出模型为ONNX格式：

import torch
# 替换为你的训练好的模型实例、对应维度的输入样例
model = YourModel()
model.load_state_dict(torch.load("your_model.pth"))
model.eval()
dummy_input = torch.randn(1, 3, 224, 224) # 按你的模型实际输入维度调整
torch.onnx.export(
    model, 
    dummy_input, 
    "converted_model.onnx", 
    export_params=True,
    opset_version=15 # 建议选稍高的兼容版本
)

2. 在Julia侧用ONNX.jl加载转换为Flux模型：

using Pkg
Pkg.add(["Flux", "ONNX"])
using Flux, ONNX

flux_model = ONNX.load("converted_model.onnx")

注意事项：

• 仅支持双方都兼容的标准算子，自定义算子需要手动写适配逻辑
• PyTorch默认张量排布为通道优先（NCHW），Flux默认是通道最后（WHCN），推理输入时需要调整维度适配

方案2：手动迁移权重（适用于含自定义算子/结构的场景） #

该方案可控性最高，不会出现算子兼容问题：

1. PyTorch侧导出所有权重为通用的.npy格式：

import numpy as np
import torch

model = YourModel()
model.load_state_dict(torch.load("your_model.pth"))
for param_name, tensor in model.state_dict().items():
    np.save(f"{param_name}.npy", tensor.detach().cpu().numpy())

2. Julia侧先手动搭建和PyTorch结构完全一致的Flux模型，再用NPZ.jl读取导出的权重文件，逐一对齐赋值给Flux模型的对应层参数即可。

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