Save PyTorch model for conversion to ONNX

Question

I'm brand new to Pytorch (and Python), I've followed this guide which trains a model and then saves the weights into a pth file: https://medium.com/@alexppppp/how-to-create-synthetic-dataset-for-computer-vision-keypoint-detection-78ba481cdafd

My understanding is that to convert a model to ONNX, you need to save the entire thing and not just the weights.

I think the relevant code is this:

for epoch in range(num_epochs):
    train_one_epoch(model, optimizer, data_loader_train, device, epoch, print_freq=1000)
    lr_scheduler.step()
    evaluate(model, data_loader_test, device)
    
# Save model weights after training
torch.save(model.state_dict(), 'keypointsrcnn_weights.pth')

Is there a simple way to save the "entire" model rather than just the weights? I've seen this in the docs but this looks like it would need to go within the epoch loop rather than after it's trained?

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

Please forgive my total lack of understanding. My intention is to try convert a PyTorch model to ONNX.

Answer 1

use torch.onnx.export. Should look something like

  arch = models.alexnet();      pic_x = 227
  dummy_input = torch.zeros((1,3, pic_x, pic_x))
  torch.onnx.export(arch, dummy_input, "alexnet.onnx", verbose=True, export_params=True, )

---

graph(%input.1 : Float(1, 3, 227, 227, strides=[154587, 51529, 227, 1], requires_grad=0, device=cpu),
      %features.0.weight : Float(64, 3, 11, 11, strides=[363, 121, 11, 1], requires_grad=1, device=cpu),
      %features.0.bias : Float(64, strides=[1], requires_grad=1, device=cpu),
      %features.3.weight : Float(192, 64, 5, 5, strides=[1600, 25, 5, 1], requires_grad=1, device=cpu),
      %features.3.bias : Float(192, strides=[1], requires_grad=1, device=cpu),
      ...
      %classifier.6.weight : Float(1000, 4096, strides=[4096, 1], requires_grad=1, device=cpu),
      %classifier.6.bias : Float(1000, strides=[1], requires_grad=1, device=cpu)):
  %17 : Float(1, 64, 56, 56, strides=[200704, 3136, 56, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[11, 11], pads=[2, 2, 2, 2], strides=[4, 4]](%input.1, %features.0.weight, %features.0.bias) # c:\python39\lib\site-packages\torch\nn\modules\conv.py:442:0
  %18 : Float(1, 64, 56, 56, strides=[200704, 3136, 56, 1], requires_grad=1, device=cpu) = onnx::Relu(%17) # c:\python39\lib\site-packages\torch\nn\functional.py:1297:0
  ...