在Unity内运行机器学习模型

1,304 阅读3分钟

因为一个项目需求需要在unity里面运行模型。由于对机器模型不了解,找过一个官方插件ML-Agent,需要配置python以及pytorch等等,还遇到py版本不合适导致报错等问题非常麻烦。后面才意识到这个插件更多用于在游戏内(根据奖励/正反馈)来训练AI行为,和项目直接给的机器模型有区别。但后面有机会还是研究一下,似乎有很大的潜力。入门可看CodeMonkey的教学视频。 后来才找到Barracuda这个插件。现在记录从0摸索使用这个插件的过程。

介绍

Introduction to Barracuda | Barracuda | 1.0.4 (unity3d.com) 大致翻译就是这个包是轻量级跨平台的神经网络库。Barracuda可以运行在GPU和CPU上运行。还能和其它AI相关的插件共用(研发中)。

安装

官方文档: Installing Barracuda | Barracuda | 1.0.4 (unity3d.com)

官方Github库:(更新更及时) barracuda-release/Documentation~/Installing.md at release/3.0.1 · Unity-Technologies/barracuda-release · GitHub

Unity 2020.3 之后在packageWindow找不到了,在manifest添加又不知道版本号。 "com.unity.barracuda" : "https://github.com/Unity-Technologies/barracuda-release.git[#<revision>]"

个人最好用的还是在PackageManager “AddByName”"com.unity.barracuda"(在不知道最新版本的情况下)。

使用方法

Getting started with Barracuda | Barracuda | 1.0.4 (unity3d.com)

处理导出后的模型

  • 从Pytorch, TensorFlow 或者 Keras 导出ONNX文件。
  • 把.onnx文件拖到Asset文件夹下,会被Unity处理成NNModel资产。 image.png
  • 脚本处理: 将NNModel资产用Model类来包装,存储形式为二进制。用一个变量来存储。
public NNModel modelAsset;
private Model m_RuntimeModel;

void Start()
{   
    m_RuntimeModel = ModelLoader.Load(modelAsset);
}    

生成worker

Barracuda核心的接口是IWorker,Worker将model分为多个可执行的任务并在CPU或GPU上运行。

var worker = WorkerFactory.CreateWorker(<WorkerFactory.Type>, m_RuntimeModel);  

处理input

算法运行可能需要输入数据,可以用单个Tensor类型的变量或者以string-Tensor为键值对的字典。

    Tensor input = new Tensor(batch, height, width, channels); 
    worker.Execute(input);
    var inputs = new Dictionary<string, Tensor>();
    inputs[name1] = new Tensor(batch, height1, width1, channels1);
    inputs[name2] = new Tensor(batch, height2, width2, channels2);
    worker.Execute(inputs);

另外顺便提一下图片数据的处理。支持Texture2DTexture2DArrayTexture3D or RenderTexture用API直接转换成Tensor,不需要逐像素处理。其中channelCount代表了tensor存照片的维度(1/3/4)。

// you can treat input pixels as 1 (grayscale), 3 (color) or 4 (color with alpha) channels
var channelCount = 3; 
var tensor = new Tensor(texture, channelCount);

关于Tensor[张量]以及其它数据处理详情见Working with data | Barracuda | 1.0.4 (unity3d.com)

worker执行input

    worker.Execute(input);

output处理

  • 单个输出,用worker.PeekOutput()
  • 多个输出,指定名字
Tensor output = m_Worker.PeekOutput(outputName);

worker.PeekOutput()这个方法只提供变量的引用,以防增加更多的内存占用;不过这也意味着变量会在下一个worker.Execute()调用或worker.Dispose()后丢失。

如果想要持有更长生命周期的tensor变量,使用worker.CopyOutput(),但不需要这个变量的时候要用Dispose()方法清除。

完整代码

(官方)

public NNModel modelAsset;
private Model m_RuntimeModel;

void Start()
{   
    m_RuntimeModel = ModelLoader.Load(modelAsset);
    m_Worker = WorkerFactory.CreateWorker(<WorkerFactory.Type>, m_RuntimeModel);
}

void Update()
{
    Tensor input = new Tensor(batch, height, width, channels); 
    m_Worker.Execute(input);
    Tensor O = m_Worker.PeekOutput("output_layer_name");
    input.Dispose();
}

拓展学习