Search Unity

  1. We are migrating the Unity Forums to Unity Discussions by the end of July. Read our announcement for more information and let us know if you have any questions.
    Dismiss Notice
  2. Dismiss Notice

Question How to use ONNX model to inference in Unity?

Discussion in 'ML-Agents' started by Dream_Surpass, Aug 29, 2023.

  1. Dream_Surpass

    Dream_Surpass

    Joined:
    Dec 2, 2022
    Posts:
    18
    In usual process, I rewrite the 'CollectObservation()' and 'OnActionReceived' function in mlagents, but I wonder what happens between these two functions.

    I want to use a pretrained onnx model to do inference by self-written code in Unity. I refer to the following code and use the barracuda to do inference. But I found that the training process which contains barracuda inference process took up a lot of GPU memory.

    What should I do? Any ideas would be appreciated!

    Code (CSharp):
    1. using System;
    2. using System.Linq;
    3. using Unity.Barracuda;
    4. using UnityEngine;
    5. public class BarracudaWorker : MonoBehaviour
    6. {
    7.     public Texture2D t;
    8.     public NNModel modelAsset;
    9.     void Start()
    10.     {
    11.         Model m_RuntimeModel = ModelLoader.Load(modelAsset);
    12.         var worker = WorkerFactory.CreateWorker(WorkerFactory.Type.ComputePrecompiled, m_RuntimeModel);
    13.         var input = new Tensor(1, 180, 180, 3);
    14.         for (var y = 0; y < 180; y++)
    15.         {
    16.             for (var x = 0; x < 180; x++)
    17.             {
    18.                 var tx = x * t.width  / 180;
    19.                 var ty = y * t.height / 180;
    20.                 /*
    21.                 input[0, 179 - y, x, 0] = t.GetPixel(tx, ty).r; //for values from 0 - 1
    22.                 input[0, 179 - y, x, 1] = t.GetPixel(tx, ty).g;
    23.                 input[0, 179 - y, x, 2] = t.GetPixel(tx, ty).b;
    24.                 */
    25.                 input[0, 179 - y, x, 0] = t.GetPixel(tx, ty).r * 255; //for values from 0 - 255
    26.                 input[0, 179 - y, x, 1] = t.GetPixel(tx, ty).g * 255;
    27.                 input[0, 179 - y, x, 2] = t.GetPixel(tx, ty).b * 255;
    28.             }
    29.         }
    30.         worker.Execute(input);
    31.         Tensor output = worker.PeekOutput("dense_1");
    32.         float[] outputBuffer = output.ToReadOnlyArray();
    33.         print(outputBuffer[0]); //print the result as a percentage
    34.         output.Dispose();
    35.     }
    36. }