This is the 10th day of my participation in the November Gwen Challenge. Check out the event details: The last Gwen Challenge 2021

An important step in pyTorch model deployment is to save the PTH model as ONNX, which is documented in this article.

Archived onnx

  • Create your own PyTorch model and load the weights
model = create_model(num_classes=2)
model.load_state_dict(load(model_path, map_location='cpu') ["model"])
Copy the code
  • Dump the onNX file
dummy_input = torch.randn(1.3.256.256, device='cpu')
torch.onnx._export(model, dummy_input, "faster_rcnn.onnx", verbose=True, opset_version=11)
Copy the code

The model is saved in the faster_rcnn.onnx file in the current directory

Verify onNX validity

  • The installationonnxruntime 库
pip install onnxruntime
Copy the code
  • Load the ONNX model and test it
import onnxruntime
from onnxruntime.datasets import get_example

def to_numpy(tensor) :
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# Test data
dummy_input = torch.randn(1.3.256.256, device='cpu')

example_model = get_example(<absolute_root_to_your_onnx_model_file>)
# netron.start(example_model) Visualize the network using the Netron Python package
sess = onnxruntime.InferenceSession(example_model)

# onnx network output
onnx_out = sess.run(None, {<input_layer_name_of_your_network>: to_numpy(dummy_input)})
print(onnx_out)

model.eval(a)with torch.no_grad():
    # PyTorch Model network output
    torch_out = model(dummy_input)
	print(torch_out)
Copy the code
  • Output:
Onnx_out [array([0., 93.246, 228.95842, 256.], [0., 2.6370468, 209.39705, 148.17822]], DType = FLOAT32), Array ([1, 1], dType =int64), array([0.1501071, 0.07568519], DType =float32)] TORCH_out [{'boxes': Tensor ([[0.0000, 93.2459, 228.9584, 256.0000], [0.0000, 2.6370, 209.3971, 148.1782]]), 'labels': Tensor ([1, 1]), 'scores' : tensor ([0.1501, 0.0757])}]Copy the code

Gets your own network input layer name

  • Sometimes, if you are not familiar with the network and do not know the name of the input layer of the model, you can use Netron to visualize your network, obtain the name of the input layer, and feed it into THE SESS of OnNX.

Attention!!

  • The pyTorch model is converted to the ONNX model using a trace-based exporter, which means that it executes by running the model once and then exporting the operators that actually participate in the operation. This also means that if your model is dynamic, for example, changing operations that depend on input data, the resulting results will be inaccurate. Similarly, a trajectory may only be valid for a specific input size (this is one reason why we need to have explicit inputs in the trajectory). We recommend checking the trace of the model to ensure that the operators being traced are reasonable

  • That is, if the network module has if.. else.. For similar branches, the ONNX model is generated by selecting one of the branches based on the initial data used, so that the resulting ONNX model retains only the structure of this branch, and the other logical capabilities in the original Pytorch model are lost in this model.

The resources

  • Pytorch.apachecn.org/docs/0.3/on…
  • www.jianshu.com/p/5a0a09fbd…