Upload 7 files

RunAutomata.cs

@@ -0,0 +1,274 @@
+using System.Collections;
+using System.Collections.Generic;
+using UnityEngine;
+using Unity.Sentis;
+using System.IO;
+using Lays = Unity.Sentis.Layers;
+
+/*
+ *      Neural Cellular Automata Inference Code
+ *      =======================================
+ *      
+ * Put this script on the Main Camera
+ * Create an image or quad in the scene.
+ * Assign an unlit transparent material to the image/quad.
+ * Draw the same material into the outputMaterial field
+ * 
+ */
+public class RunAutomata : MonoBehaviour
+{
+    //Change this to load a different model:
+    public AutomataNames automataName = AutomataNames.Poop;
+
+    //Reduce this to make it run slower
+    [Range(0f, 1f)]
+    public float stepSize = 1.0f;
+
+    const BackendType backend = BackendType.GPUCompute;
+
+    //Drag your unlit transparent material here for drawing the output
+    public Material outputMaterial;
+
+    //optional material for average alpha 
+    public Material avgAlphaMaterial;
+
+    public enum AutomataNames { Lizard, Turtle ,Poop};
+   
+    //Model parameters
+    const int trainedResolution = 40;
+    const int trainedPool = 16;
+    const int alphaBlocks = 4;
+    int m_paddedImageSize;
+    int m_trainedHiddenStates;
+
+    //Workers to run the networks
+    private IWorker m_WorkerStateUpdate;
+    private IWorker m_WorkerClip;
+
+    private TensorFloat m_currentStateTensor;
+    private RenderTexture m_currentStateTexture;
+    private RenderTexture m_currentBlockAlphaStateTexture;
+
+    Ops m_ops;
+    ITensorAllocator m_allocator;
+
+    void Start()
+    {
+        m_allocator = new TensorCachingAllocator();
+        m_ops = WorkerFactory.CreateOps(backend, m_allocator);
+
+        Application.targetFrameRate = 60;
+
+        LoadAutomataModel();
+
+        CreateProcessingModel();
+
+        SetupState();
+
+        SetupTextures();
+
+        DrawDotAt(m_paddedImageSize / 2, m_paddedImageSize / 2);
+    }
+    void LoadAutomataModel() {
+
+        Model m_ModelStateUpdate = null;
+
+        switch (automataName) {
+            case AutomataNames.Lizard:
+                m_ModelStateUpdate = ModelLoader.Load(Application.streamingAssetsPath + "/lizard.sentis");
+                break;
+            case AutomataNames.Turtle:
+                m_ModelStateUpdate = ModelLoader.Load(Application.streamingAssetsPath + "/turtle.sentis");
+                break;
+            case AutomataNames.Poop:
+                m_ModelStateUpdate = ModelLoader.Load(Application.streamingAssetsPath + "/poop.sentis");
+                break;
+        }
+        m_trainedHiddenStates = m_ModelStateUpdate.inputs[0].shape[3].value;
+
+        m_paddedImageSize = trainedResolution + trainedPool * 2;
+
+        m_WorkerStateUpdate = WorkerFactory.CreateWorker(backend, m_ModelStateUpdate, false);
+
+    }
+    void CreateProcessingModel() { 
+
+        var m_Model = new Model();
+
+        var input0 = new Model.Input
+        {
+            name = "input0",
+            shape = (new SymbolicTensorShape(1, m_trainedHiddenStates, m_paddedImageSize, m_paddedImageSize)),
+            dataType=DataType.Float
+        };
+
+        var input1 = new Model.Input
+        {
+            name = "input1",
+            shape = (new SymbolicTensorShape(1, m_trainedHiddenStates, m_paddedImageSize, m_paddedImageSize)),
+            dataType = DataType.Float
+        };
+
+        var inputStepSize = new Model.Input
+        {
+            name = "inputStepSize",
+            shape = new SymbolicTensorShape(1, 1, 1, 1),
+            dataType = DataType.Float
+        };
+
+        m_Model.inputs.Add(input0);
+        m_Model.inputs.Add(input1);
+        m_Model.inputs.Add(inputStepSize);
+
+        m_Model.AddConstant(new Lays.Constant("aliveRate", new TensorFloat(new TensorShape(1, 1, 1, 1), new[] { 0.1f })));
+
+        m_Model.AddConstant(new Lays.Constant("sliceStarts", new int[] { 0, 3, 0, 0 }));
+        m_Model.AddConstant(new Lays.Constant("sliceEnds", new[] { 1, 4 ,m_paddedImageSize, m_paddedImageSize }));
+
+        m_Model.AddLayer(new Lays.Slice("sliceI0", "input0", "sliceStarts", "sliceEnds"));
+        m_Model.AddLayer(new Lays.MaxPool("maxpool0", "sliceI0", new[] { 3, 3 }, new[] { 1, 1 }, new[] { 1, 1, 1, 1 }));
+        m_Model.AddLayer(new Lays.Greater("pre_life_mask", "maxpool0", "aliveRate")); //INT
+        
+        m_Model.AddLayer(new Lays.Mul("input1_stepsize",  "input1", "inputStepSize" ));
+        
+        m_Model.AddLayer(new Lays.RandomUniform("random", new int[] { 1, 1, m_paddedImageSize, m_paddedImageSize}, 0.0f, 1.0f, 0));
+        m_Model.AddConstant(new Lays.Constant("fireRate", new TensorFloat(new TensorShape(1, 1, 1, 1), new[] { 0.5f })));
+        m_Model.AddLayer(new Lays.LessOrEqual("lessEqualFireRateINT", "random", "fireRate"));
+
+        m_Model.AddLayer(new Lays.Cast("lessEqualFireRate", "lessEqualFireRateINT", DataType.Float));
+        
+        m_Model.AddLayer(new Lays.Mul("mul", "input1_stepsize", "lessEqualFireRate" ));
+        
+        m_Model.AddLayer(new Lays.Add("add", "input0", "mul" ));
+        
+        m_Model.AddLayer(new Lays.Slice("sliceI1", "add", "sliceStarts", "sliceEnds"));
+        m_Model.AddLayer(new Lays.MaxPool("maxpool1", "sliceI1", new [] { 3 ,3 }, new[] { 1, 1 }, new[] {1, 1, 1, 1}));
+        m_Model.AddLayer(new Lays.Greater("post_life_mask", "maxpool1", "aliveRate"));
+
+        
+        m_Model.AddLayer(new Lays.And("andINT", "pre_life_mask", "post_life_mask"));
+        m_Model.AddLayer(new Lays.Cast("and", "andINT", DataType.Float));
+        
+        m_Model.AddLayer(new Lays.Mul("outputState", "add", "and" ));
+
+        m_Model.AddConstant(new Lays.Constant("sliceStarts2", new[] { 0, 0, trainedPool, trainedPool }));
+        m_Model.AddConstant(new Lays.Constant("sliceEnds2", new[] { 1, 4, m_paddedImageSize - trainedPool, m_paddedImageSize - trainedPool }));
+
+        m_Model.AddLayer(new Lays.Slice("outputImage", "outputState", "sliceStarts2", "sliceEnds2"));
+        
+        m_Model.AddLayer(new Lays.Slice("outputIC", "outputImage", "sliceStarts", "sliceEnds"));
+
+        int blockSize = trainedResolution / alphaBlocks;
+        m_Model.AddLayer(new Lays.AveragePool("avgPoolBlocks", "outputIC", new[] { blockSize, blockSize }, new[] { blockSize, blockSize }, new[] { 1, 1, 1, 1 }));
+
+        m_Model.outputs.Add("outputState");
+        m_Model.outputs.Add("outputImage");
+        m_Model.outputs.Add("avgPoolBlocks");
+        
+        m_WorkerClip = WorkerFactory.CreateWorker(BackendType.GPUCompute, m_Model);
+
+    }
+
+    void SetupState()
+    {
+        float[] data = new float[1 * m_paddedImageSize * m_paddedImageSize * m_trainedHiddenStates];
+        m_currentStateTensor = new TensorFloat(new TensorShape(1, m_trainedHiddenStates, m_paddedImageSize, m_paddedImageSize), data);
+    }
+
+    void SetupTextures()
+    {
+        m_currentStateTexture = new RenderTexture(trainedResolution, trainedResolution, 0)
+        {
+            enableRandomWrite = true
+        };
+        outputMaterial.mainTexture = m_currentStateTexture;
+
+        if (avgAlphaMaterial)
+        {
+            m_currentBlockAlphaStateTexture = new RenderTexture(alphaBlocks, alphaBlocks, 0)
+            {
+                enableRandomWrite = true
+            };
+            outputMaterial.mainTexture = m_currentBlockAlphaStateTexture;
+        }
+    }
+
+    void DrawDotAt(int x,int y)
+    {
+        m_currentStateTensor.MakeReadable();
+
+        float[] data = m_currentStateTensor.ToReadOnlyArray();
+        for (int k = 3; k < 16; k++)
+        {
+            data[m_paddedImageSize * m_paddedImageSize * k + m_paddedImageSize * y + x] = 1f;
+        }
+        Replace(ref m_currentStateTensor, new TensorFloat(m_currentStateTensor.shape, data));
+    }
+
+    void Update()
+    {
+        DoInference();
+        if (Input.GetKeyDown(KeyCode.Escape))
+        {
+            Application.Quit();
+        }
+        if (Input.GetKeyDown(KeyCode.Space))
+        {
+            DrawDotAt(UnityEngine.Random.Range(0, m_paddedImageSize), UnityEngine.Random.Range(0, m_paddedImageSize));
+        }
+    }
+
+    void Replace(ref TensorFloat A, TensorFloat B)
+    {
+        A?.Dispose();
+        A = B;
+    }
+
+    void DoInference() {
+
+        using var stepSizeTensor = new TensorFloat(new TensorShape(1, 1, 1, 1), new float[] { stepSize });
+
+        using var currentStateTensorT = m_ops.Transpose(m_currentStateTensor, new int[] { 0, 2, 3, 1 });
+
+        m_WorkerStateUpdate.Execute(currentStateTensorT);
+        TensorFloat outputStateT = m_WorkerStateUpdate.PeekOutput() as TensorFloat;
+
+        using var outputState = m_ops.Transpose(outputStateT, new int[] { 0, 3, 1, 2 });
+
+        var inputs = new Dictionary<string, Tensor>() { 
+            { "input0", m_currentStateTensor }, //float
+            { "input1", outputState }, //float
+            { "inputStepSize", stepSizeTensor }  //float
+        };
+        m_WorkerClip.Execute(inputs);
+
+        TensorFloat clippedState = m_WorkerClip.PeekOutput("outputState") as TensorFloat;
+        TensorFloat outputImage = m_WorkerClip.PeekOutput("outputImage") as TensorFloat;
+        TensorFloat blockAvgAlphaState = m_WorkerClip.PeekOutput("avgPoolBlocks") as TensorFloat;
+
+        if (m_currentStateTexture)
+        {
+            TextureConverter.RenderToTexture(outputImage, m_currentStateTexture);
+        }
+
+        if (m_currentBlockAlphaStateTexture)
+        {
+            TextureConverter.RenderToTexture(blockAvgAlphaState, m_currentBlockAlphaStateTexture);
+        }
+
+        Replace(ref m_currentStateTensor, clippedState);
+        m_currentStateTensor.TakeOwnership();
+    }
+
+    void OnDestroy()
+    {
+        m_currentStateTensor.Dispose();
+
+        m_WorkerStateUpdate.Dispose();
+        m_WorkerClip.Dispose();
+
+        m_ops?.Dispose();
+        m_allocator?.Dispose();
+    }
+   
+}

lizard.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfc3f2761a6c13c3985c5c2d4afdb8495c63916be0166708b6c8de5495301af9
+size 62332

poop.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3178f88f94defda03b83b33b8a9c6326edc16dfd496ad16e2705a5ac86be35ef
+size 62330

turtle.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d64f01dd80339ec69bebe2db259cad36ed6d2f1fe7c5d5c12d3e309a6e877ee2
+size 75209