A free, anonymous, minimalist and open source paste tool.

15-06-2021 22:28:16
model_input_name = "mobilefacenet_model_best.pth.tar"
model_output_name = "mobilefacenet_model_best.pth.tar"

import torch
import torch.onnx
import torchvision
import torchvision.models as models
import sys

onnx_model_path = ""

# https://pytorch.org/hub/pytorch_vision_densenet/
model = torch.load(model_input_name,map_location=torch.device('cpu'),pretrained=True)

# set the model to inference mode

# Create some sample input in the shape this model expects
# This is needed because the convertion forward pass the network once
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "new_model.onnx", verbose=True)