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

Question

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

Answer 1

我唯一能想到的方法就是

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")

1. 使用.onnx使用ONNX.jl加载模型。看起来这个库目前正在重新构建中，但是旧的API可能对您有用。再检查一遍，似乎在加载模型后可能会有退货空间。

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