Run PyTorch inside Unity

zeroize: <p>Hi! I’m working on a unity project that need to run my trained model inside the app. I tried to convert the .pt file into .onnx so it can be run inside the unity (using barracuda), however the prediction is changed so I couldn’t get the correct data.</p>
<p>This is the code that I used to export the pt file:</p>
<p>import torch<br>
import torch.nn as nn<br>
import torchvision.models as models<br>
import numpy as np<br>
import random<br>
import onnxruntime as ort<br>
import json</p>
<h1><a name="set-random-seeds-for-reproducibility-1" class="anchor" href="#set-random-seeds-for-reproducibility-1"></a>Set random seeds for reproducibility</h1>
<h1><a name="torchmanual_seed42-2" class="anchor" href="#torchmanual_seed42-2"></a>torch.manual_seed(42)</h1>
<h1><a name="nprandomseed42-3" class="anchor" href="#nprandomseed42-3"></a>np.random.seed(42)</h1>
<h1><a name="randomseed42-4" class="anchor" href="#randomseed42-4"></a>random.seed(42)</h1>
<p>class MyResNet18Model(nn.Module):<br>
def <strong>init</strong>(self, num_classes=10):<br>
super(MyResNet18Model, self).<strong>init</strong>()<br>
self.resnet18 = models.resnet18(pretrained=False)<br>
self.resnet18.fc = nn.Linear(self.resnet18.fc.in_features, num_classes)</p>
<pre><code>def forward(self, x):
    return self.resnet18(x)
</code></pre>
<p>model = MyResNet18Model(num_classes=10)</p>
<h1><a name="load-checkpoint-5" class="anchor" href="#load-checkpoint-5"></a>Load checkpoint</h1>
<p>checkpoint = torch.load(‘./Training Model/pf_model_copy.pt’)<br>
model.load_state_dict(checkpoint[‘state_dict’], strict=False)</p>
<h1><a name="set-model-to-evaluation-mode-6" class="anchor" href="#set-model-to-evaluation-mode-6"></a>Set model to evaluation mode</h1>
<p>model.eval()</p>
<h1><a name="create-dummy-input-7" class="anchor" href="#create-dummy-input-7"></a>Create dummy input</h1>
<p>dummy_input = torch.randn(1, 3, 720, 720)</p>
<h1><a name="get-pytorch-model-output-8" class="anchor" href="#get-pytorch-model-output-8"></a>Get PyTorch model output</h1>
<p>with torch.no_grad():<br>
pytorch_output = model(dummy_input)</p>
<h1><a name="export-model-to-onnx-format-9" class="anchor" href="#export-model-to-onnx-format-9"></a>Export model to ONNX format</h1>
<p>torch.onnx.export(<br>
model,  # model being run<br>
dummy_input,  # model input (or a tuple for multiple inputs)<br>
“./Training Model/pf_model_copy_v4.onnx”,  # where to save the model (can be a file or file-like object)<br>
export_params=True,  # store the trained parameter weights inside the model file<br>
opset_version=11,  # the ONNX version to export the model to<br>
do_constant_folding=True,  # whether to execute constant folding for optimization<br>
input_names=[‘input’],  # the model’s input names<br>
output_names=[‘output’],  # the model’s output names<br>
dynamic_axes={‘input’: {0: ‘batch_size’},  # variable length axes<br>
‘output’: {0: ‘batch_size’}}<br>
)</p>
<p>print(“Model has been converted to ONNX format and saved as pf_model_copy_v4.onnx”)</p>
<h1><a name="create-onnx-runtime-session-10" class="anchor" href="#create-onnx-runtime-session-10"></a>Create ONNX runtime session</h1>
<p>ort_session = ort.InferenceSession(“./Training Model/pf_model_copy_v4.onnx”)</p>
<h1><a name="function-to-convert-a-pytorch-tensor-to-a-numpy-array-11" class="anchor" href="#function-to-convert-a-pytorch-tensor-to-a-numpy-array-11"></a>Function to convert a PyTorch tensor to a NumPy array</h1>
<p>def to_numpy(tensor):<br>
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()</p>
<h1><a name="prepare-inputs-for-onnx-runtime-12" class="anchor" href="#prepare-inputs-for-onnx-runtime-12"></a>Prepare inputs for ONNX runtime</h1>
<p>ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(dummy_input)}</p>
<h1><a name="run-inference-13" class="anchor" href="#run-inference-13"></a>Run inference</h1>
<p>ort_outs = ort_session.run(None, ort_inputs)</p>
<h1><a name="print-outputs-14" class="anchor" href="#print-outputs-14"></a>Print outputs</h1>
<p>print(“PyTorch output:”, pytorch_output)<br>
print(“ONNX output:”, ort_outs[0])</p>
<h1><a name="save-labels-to-json-15" class="anchor" href="#save-labels-to-json-15"></a>Save labels to JSON</h1>
<p>labels = {<br>
0: ‘Grafitti_Gorilla’,<br>
1: ‘Hydro_Phantasm’,<br>
2: ‘Purple_Fighter’,<br>
3: ‘Pyro_Monkey’,<br>
4: ‘Red_panda’,<br>
5: ‘Smoke_Face’,<br>
6: ‘Split_Girl’,<br>
7: ‘Swap_Kitty’,<br>
8: ‘Taiyaki’,<br>
9: ‘UFO_Cat’<br>
}</p>
<p>with open(‘./Training Model/labels.json’, ‘w’) as f:<br>
json.dump(labels, f)</p>
<p>Does anybody can help? Thank you!</p>

Friedrich: <p>Hello,</p>
<p>I do not know if this helps, but I could export my .pth file from my feed forward neural net to .onnx with:</p>
<pre><code>def multi_layer_perceptron_export_onnx(self, input_size):

    model_loaded = MLP(self.input_size_mlp, self.hidden_size_1, self.hidden_size_2, self.hidden_size_3)
    model_loaded.load_state_dict(
        torch.load(path_to_project_for_stored_models + 'multi_layer_perceptron_classification_model.pth'))

    # Set the model to evaluation mode
    model_loaded.eval()

    # Provide an example input that represents the input size of your model
    # Usually, it's a tensor with the same shape as your input data

    # Sliding Window will have the length 300
    example_input = torch.randn(input_size, 1)

    # Convert the PyTorch model to ONNX format
    onnx_path = path_to_project_for_stored_models + "multi_layer_perceptron_classification_model.onnx"
    torch.onnx.export(model_loaded, example_input, onnx_path, verbose=True, input_names=["input"],
                      output_names=["output"], opset_version=15, export_params=True)
</code></pre>
<p>I played a bit around and figured out, that max Version: 15 worked in combination with export_params=True to use the model in Unity.<br>
Let me know if this helps.</p>
<p>Kind regards</p>