<jupyter_start><jupyter_text>Import a model from ONNX and run using PyTorch We demonstrate how to import a model from ONNX and convert to PyTorch  Imports<jupyter_code>import os
import operator as op
import warnings; warnings.simplefilter(action='ignore', category=FutureWarning)


import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from torch.autograd import Variable
import onnx

import gamma
from gamma import convert, protobuf, utils<jupyter_output><empty_output><jupyter_text>1: Download the model<jupyter_code>fpath = utils.get_file('https://s3.amazonaws.com/download.onnx/models/squeezenet.tar.gz')
onnx_model = onnx.load(os.path.join(fpath, 'squeezenet/model.onnx'))

inputs = [i.name for i in onnx_model.graph.input if 
          i.name not in {x.name for x in onnx_model.graph.initializer}]
outputs = [o.name for o in onnx_model.graph.output]<jupyter_output><empty_output><jupyter_text>2: Import into Gamma<jupyter_code>graph = convert.from_onnx(onnx_model)
constants = {k for k, (v, i) in graph.items() if v['type'] == 'Constant'}
utils.draw(gamma.strip(graph, constants))<jupyter_output><empty_output><jupyter_text>3: Convert to PyTorch<jupyter_code>make_node = gamma.make_node_attr
def torch_padding(params):
    padding = params.get('pads', [0,0,0,0])
    assert (padding[0] == padding[2]) and (padding[1] == padding[3]) 
    return (padding[0], padding[1])
    
torch_ops = {
    'Add': lambda params: op.add,
    'Concat': lambda params: (lambda *xs: torch.cat(xs, dim=params['axis'])),
    'Constant': lambda params: nn.Parameter(torch.FloatTensor(params['value'])), 
    'Dropout': lambda params: nn.Dropout(params.get('ratio', 0.5)).eval(), #.eval() sets to inference mode. where should this logic live?
    'GlobalAveragePool': lambda params: nn.AdaptiveAvgPool2d(1),
    'MaxPool': lambda params: nn.MaxPool2d(params['kernel_shape'], stride=params.get('strides', [1,1]),
                                            padding=torch_padding(params), 
                                            dilation=params.get('dilations', [1,1])),
    'Mul': lambda params: op.mul,
    'Relu': lambda params: nn.ReLU(),
    'Softmax': lambda params: nn.Softmax(dim=params.get('axis', 1)), 
}    

def torch_op(node, inputs):
    if node['type'] in torch_ops:
        op = torch_ops[node['type']](node['params'])
        return make_node('Torch_op', {'op': op}, inputs)
    return (node, inputs)


def torch_conv_node(params, x, w, b):
    ko, ki, kh, kw = w.shape
    group = params.get('group', 1)
    ki *= group
    conv = nn.Conv2d(ki, ko, (kh,kw), 
                     stride=tuple(params.get('strides', [1,1])),
                     padding=torch_padding(params), 
                     dilation=tuple(params.get('dilations', [1,1])),
                     groups=group)
    conv.weight = nn.Parameter(torch.FloatTensor(w))
    conv.bias = nn.Parameter(torch.FloatTensor(b))
    return make_node('Torch_op', {'op': conv}, [x])

def convert_to_torch(graph):
    v, _ = gamma.var, gamma.Wildcard
    conv_pattern = {
        v('conv'): make_node('Conv', v('params'), [v('x'), v('w'), v('b')]),
        v('w'): make_node('Constant', {'value': v('w_val')}, []),
        v('b'): make_node('Constant', {'value': v('b_val')}, [])
    }
    matches = gamma.search(conv_pattern, graph)
    g = gamma.union(graph, {m[v('conv')]: 
            torch_conv_node(m[v('params')], m[v('x')], m[v('w_val')], m[v('b_val')]) 
            for m in matches})
    remove = {m[x] for m in matches for x in (v('w'), v('b'))}
    g = {k: torch_op(v, i) for k, (v, i) in g.items() if k not in remove}
    return g

def torch_graph(graph):
    return gamma.FuncCache(lambda k: graph[k][0]['params']['op'](*[tg[x] for x in graph[k][1]]))
g = convert_to_torch(graph)
utils.draw(g)<jupyter_output><empty_output><jupyter_text>4: Load test example and check PyTorch output<jupyter_code>def load_onnx_tensor(fname):
    tensor = onnx.TensorProto()
    with open(fname, 'rb') as f:
        tensor.ParseFromString(f.read())
    return protobuf.unwrap(tensor)

input_0 = load_onnx_tensor(os.path.join(fpath, 'squeezenet/test_data_set_0/input_0.pb'))
output_0 = load_onnx_tensor(os.path.join(fpath, 'squeezenet/test_data_set_0/output_0.pb'))
tg = torch_graph(g)
tg[inputs[0]] = Variable(torch.Tensor(input_0))
torch_outputs = tg[outputs[0]]

np.testing.assert_almost_equal(output_0, torch_outputs.data.numpy(), decimal=5)
print('Success!')<jupyter_output>Success!
{'fraction_non_alphanumeric': 0.16223463687150838, 'fraction_numerical': 0.009832402234636871, 'mean_word_length': 4.046223224351747, 'pattern_counts': {'":': 0, '<': 20, '<?xml version=': 0, '>': 20, 'https://': 1, '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': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '9247', 'license': 'MIT', 'n_tokens_mistral': 1600, 'n_tokens_neox': 1572, 'n_words': 342, 'path': 'examples/onnx_to_pytorch.ipynb', 'repo_name': 'samgd/gamma', 'size': '4475'}