How to use a .pth model in Flux.jl?

user_7619808 0
Viewed 140

I have a model trained in PyTorch, saved in .pth format. Is it possible to use and load that model in Flux.jl? I looked around but did not see this mentioned anywhere in the Flux docs.

flux.jl julia
1 Answers

The only way I can think of is

  1. Convert .pth to .onnx by
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. Load the model from .onnxusing ONNX.jl. It seems that this library is currently under re-construction but old API might work for you. Double check it, it seems that there could be descrepancies after loading the model.

Also, this discussion is relevant: https://github.com/FluxML/ML-Coordination-Tracker/issues/10

user_141860870
Related