因为一个项目需求需要在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资产。
- 脚本处理: 将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);
另外顺便提一下图片数据的处理。支持Texture2D, Texture2DArray, Texture3D 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();
}