<jupyter_start><jupyter_text>Pytorch to ONNX conversionThis notebook shows how to convert your trained Pytorch model to ONNX, the generic format supported by DIANNA.Based on tutorial at https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html<jupyter_code>import numpy as np

import torch
import torch.nn as nn
import torch.nn.init as init
from torch.utils import model_zoo

import onnx
import onnxruntime as ort<jupyter_output><empty_output><jupyter_text>Create an example model<jupyter_code>class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()

        self.relu = nn.ReLU(inplace=inplace)
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

        self._initialize_weights()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.pixel_shuffle(self.conv4(x))
        return x

    def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv4.weight)

# Create the super-resolution model by using the above model definition.
torch_model = SuperResolutionNet(upscale_factor=3)<jupyter_output><empty_output><jupyter_text>Instead of training, apply pre-determined weights<jupyter_code># get existing weights
# Load pretrained model weights
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'

# Initialize model with the pretrained weights
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=torch.device('cpu')))

# set the model to inference mode
torch_model.eval()<jupyter_output><empty_output><jupyter_text>Evaluate the model on some random input<jupyter_code># Input to the model
x = torch.randn(1, 1, 224, 224, requires_grad=True)
pred = torch_model(x)<jupyter_output><empty_output><jupyter_text>Export the model in ONNX format<jupyter_code>onnx_file = 'pytorch_super_resolution_net.onnx'
# Export the model
torch.onnx.export(torch_model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  onnx_file,                 # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes
                                'output' : {0 : 'batch_size'}})<jupyter_output><empty_output><jupyter_text>Verify the PyTorch and ONNX predictions match<jupyter_code># verify the ONNX model is valid
onnx_model = onnx.load(onnx_file)
onnx.checker.check_model(onnx_model)

# get ONNX predictions
sess = ort.InferenceSession(onnx_file)
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

onnx_input = {input_name: x.detach().numpy().astype(np.float32)}
pred_onnx = sess.run([output_name], onnx_input)[0]

# compare to Pytorch predictions
np.allclose(pred.detach().numpy(), pred_onnx, atol=1e-5)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10904458598726115, 'fraction_numerical': 0.02140127388535032, 'mean_word_length': 3.640661938534279, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '496638', 'license': 'Apache-2.0', 'n_tokens_mistral': 1361, 'n_tokens_neox': 1299, 'n_words': 301, 'path': 'onnx_conversion_scripts/pytorch.ipynb', 'repo_name': 'dianna-ai/dianna-exploration', 'size': '3925'}