Memory leak in Conv operator?

I am getting a memory leak from the ops.Conv operator. Here is my code:

using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using Unity.Sentis;

public class Waves : MonoBehaviour
{
    const int W = 256;
    TensorFloat waves, wavesVel, kernel, bias;
    public GameObject quad;
    static Ops ops;

    const float dt = 0.01f;
    ITensorAllocator m_Allocator;

    RenderTexture renderTexture;

    int[] stride = new int[] { 1, 1 };
    int[] pad = new int[] { 1, 1, 1, 1 };
    int[] dilation = new int[] { 1, 1 };
    // Start is called before the first frame update
    void Start()
    {
        m_Allocator = new TensorCachingAllocator();
        ops = WorkerFactory.CreateOps(BackendType.GPUCompute, m_Allocator);

        float[] input = new float[W * W];
        input[W * (W / 2) + W / 2] = 100f;
        waves = new TensorFloat(new TensorShape(1, 1, W, W), input);
        wavesVel = new TensorFloat(new TensorShape(1, 1, W, W), input);
        kernel = new TensorFloat(new TensorShape(1, 1, 3, 3), new float[]
        {
            0,  1,  0,
            1, -4,  1,
            0,  1,  0
        });
        bias = new TensorFloat(new TensorShape(1, 1, 1, 1), new float[] { 0 });
        renderTexture = new RenderTexture(W, W, 1);
        quad.GetComponent<Renderer>().material.mainTexture = renderTexture;
    }

    void Update()
    {
        for (int N = 0; N < 10; N++)
        {
            DoIteration();
        }

        TextureConverter.RenderToTexture(waves, renderTexture);

        if (Input.GetKeyDown(KeyCode.Space))
        {
            waves.MakeReadable();
            float[] values = waves.ToReadOnlyArray();
            int X = UnityEngine.Random.Range(0, W);
            int Y = UnityEngine.Random.Range(0, W);
            values[W * Y + X] = 100f;
            Replace(ref waves, new TensorFloat(new TensorShape(1, 1, W, W), values));
        }
    }

    void DoIteration()
    {
        using TensorFloat DD = ops.Conv(waves, kernel, bias, 1, stride, pad, dilation); //<--memory leak
        using TensorFloat dwavesVel = ops.Mul(DD, dt);
        Replace(ref wavesVel, ops.Add(wavesVel, dwavesVel));
        using TensorFloat dwaves = ops.Mul(wavesVel, dt);
        Replace(ref waves, ops.Add(waves, dwaves));
    }
    static void Replace(ref TensorFloat A, TensorFloat B)
    {
        A?.Dispose();
        A = B;
    }

    private void OnApplicationQuit()
    {
        kernel?.Dispose();
        bias?.Dispose();
        waves?.Dispose();
        wavesVel?.Dispose();
        ops?.Dispose();
        m_Allocator?.Dispose();
    }
}

Unity 2023.2. Sentis 1.3.0-pre.2

3 posts - 2 participants

Read full topic