如何在 Flux.jl 中使用 .pth 模型？

【问题标题】：How to use a .pth model in Flux.jl?如何在 Flux.jl 中使用 .pth 模型？
【发布时间】：2021-09-18 03:32:54
【问题描述】：

我有一个用 PyTorch 训练的模型，以 .pth 格式保存。是否可以在 Flux.jl 中使用和加载该模型？我环顾四周，但在 Flux 文档中的任何地方都没有提到这一点。

【问题讨论】：

标签： julia flux.jl

【解决方案1】：

我能想到的唯一方法是

通过

.pth

.onnx

import torch.onnx
import torchvision
import torch

dummy_input = #...
model = #...
state_dict = torch.load('model.pth')
model.load_state_dict(state_dict)
torch.onnx.export(model, dummy_input, "model.onnx")

使用ONNX.jl从.onnx加载模型。看来这个库目前正在重建中，但旧的 API 可能对你有用。仔细检查一下，加载模型后似乎存在差异。

另外，这个讨论是相关的：https://github.com/FluxML/ML-Coordination-Tracker/issues/10

【讨论】：