Netron 是一个支持 Tensorflow ，PyTorch ，MXNet ，NCNN ， PaddlePaddle 等深度模型格式的可视化框架。去年国庆前的时候我稍微研究了下相关的代码，重点关注其将其是如何设计出一套兼容不同模型格式表征，用来归一化展现不同的深度学习框架模型。
研究完成后，我利用如下两个 Commit 作为 Pull Request 提交给了作者，用以支持 MNN 的模型可视化。
从中也不难看出我扎实的英语表述能力（我果然是个国际化人才）。
这篇文章会从架构设计、标准定义、巧用JS解析等几个方面来阐述
架构设计
整体上，按照我个人的理解，Netron 的架构可以简要展现如下：
最基础的应用部分及运行环境，是 Electron 这个跨平台框架直接呈现的。
当然，一些诸如基础zip/gzip用于解压等等的库我们也统一归类到支撑里。
然后是一套经典的 MVC 的结构，app.js 作为整体的 controller ，负责整个应用的功能逻辑，如导出图片、菜单管理、保存加载等等。这一层我们需要的做事非常少，只要将 MNN 支持的模型后缀 .mnn 注册进去即可。 然后是是对应的 view.js，这块实际上还是一层 controller，类比我们常说的子控制器，专门用于处理主视图的逻辑，如下图所示：
从这块开始，我们就要注意了，因为这里开始通过工厂方法对应的根据读取文件类型的不同，托管给了不同的自定义 xxx.js 来处理后续步骤。 比如.mar，model，prototxt 等格式的模型会首先托管给 mxnet.js来处理。如果存在重名，则按照先后顺序依次尝试。
view.ModelFactoryService = class {
constructor(host) {
this._host = host;
this._extensions = [];
this.register('./onnx', [ '.onnx', '.pb', '.pbtxt', '.prototxt' ]);
this.register('./mxnet', [ '.mar', '.model', '.json', '.params' ]);
this.register('./keras', [ '.h5', '.hd5', '.hdf5', '.keras', '.json', '.model' ]);
this.register('./coreml', [ '.mlmodel' ]);
this.register('./caffe', [ '.caffemodel', '.pbtxt', '.prototxt', '.pt' ]);
this.register('./caffe2', [ '.pb', '.pbtxt', '.prototxt' ]);
this.register('./pytorch', [ '.pt', '.pth', '.pkl', '.h5', '.t7', '.model', '.dms', '.pth.tar', '.ckpt', '.bin' ]);
this.register('./torch', [ '.t7' ]);
this.register('./torchscript', [ '.pt', '.pth' ]);
this.register('./mnn', ['.mnn', '.tflite']);
this.register('./tflite', [ '.tflite', '.lite', '.tfl', '.bin' ]);
this.register('./tf', [ '.pb', '.meta', '.pbtxt', '.prototxt', '.json' ]);
this.register('./sklearn', [ '.pkl', '.joblib', '.model' ]);
this.register('./cntk', [ '.model', '.cntk', '.cmf', '.dnn' ]);
this.register('./openvino', [ '.xml' ]);
this.register('./darknet', [ '.cfg' ]);
this.register('./paddle', [ '.paddle', '__model__' ]);
this.register('./ncnn', [ '.param', '.bin', '.cfg.ncnn', '.weights.ncnn']);
this.register('./dl4j', [ '.zip' ]);
this.register('./mlnet', [ '.zip']);
}
在这上层是一层标准定义层，用于抹平不同模型之间的表达方式，用归一化的逻辑来进行处理，至于怎么把自己的模型表征映射成归一化的逻辑，就需要编写对应 xxx.js 来自行处理，后文会以 MNN 来进行举例。
最上层就是对应各个深度框架自行的逻辑处理了。其中包含了数据格式及对应解析（如 flatbuffer）、内容校验、构图等等，后文也会用 MNN 举例说明。
标准定义
这一环是一个很不起眼但是却非常重要的环节。 每种深度模型框架都有其自定义的模块结构和模块构成，一般都以 Flatbuffer Schema 的形式构成。（当然也有例外）以MNN 为例，其对应的模型结构大致如下图所示：
上图引用自FlatBuffers，MNN模型存储结构基础 —- 无法解读MNN模型文件的秘密
同理， TFLite 的模型也可见 TFLite.schema，不再赘述。
从定义中不难看出，TFLite 有 model，graph ，SubGraph 等；而 MNN 对应的就是Net；再往下一层 TFLite 有 Operator 和 Options；而 MNN 有 OP 和OPParameter；至于 NCNN 则是 Layer。
如果是从整个架构角度去兼容不同的框架，必然会有着大量的 messy code。因此作者定义了一套标准表征，让不同的深度模型自己去解析，然后附着自身的逻辑到这同一套表征上。
-
Model ，表示模型的静态表示。
-
Graph ，表示模型的计算图表示。
-
Node ，一个操作对应一个节点。
-
Tensor ，输入输出数据。
-
Parameter ，对应的属性。
-
Argument ，对应的属性值。
上述 Parameter 和 Argument可以简单认为一一对应吧，都认为是属性值即可。
一图胜千言，下图比较好的展现了术语和对应的表征：
这样不同的框架模型只要在自己对应的 xxx.js 中，把图，OP 层对应的数据填充至对应的地方即可。
这里依然以 MNN 举例：
- 我们不存在
subgraph 的概念，直接把 Model 和 Graph 等价于一个 net即可。
- 从
net 中取出 oplist ，对应创建成 Node。
- 从
oplist 中每个 op ，取出对应的 tensorIndex，根据 net 的 tensorName 和tensorIndex 来创建对应的 tensor 。
- 从
op 中根据 opparameter 的种类，从 op.main 中取出不同的数据来填入 paramter / argument，这块是解析的大头，如果没想好方式，就会非常浪费时间，下文重点说。
数据格式
诸如 MNN ，TFlite 都选用了 Flatbuffer 来进行数据的保存，而官方的 flatc 程序支持直接根据定义的 schema 文件生成对应的 generated.js，命令如下：
./flatc -s ~/yourPathTo/MNN/schema/default/Type.fbs
这个我看了下很多的同学的在处理多 Schema 定义的时候是对应的一个个生成 generated.js，这样维护成本比较大，既然我们的已经使用了 include 机制，我们直接在生成过程中合并即可，如下所示：
./flatc --js -I ~/yourPathTo/MNN/schema/default/ ~/yourPathTo/MNN/schema/default/MNN.fbs --gen-all
这里有两个参数注意下：
-
-I，表示 include 从哪个路径进行搜索。
-
--gen-all，表示自动对生成的所有文件合并。
生成代码大致如下：
/**
* @param {number} i
* @param {flatbuffers.ByteBuffer} bb
* @returns {MNN.Blob}
*/
MNN.Blob.prototype.__init = function(i, bb) {
this.bb_pos = i;
this.bb = bb;
return this;
};
/**
* @param {flatbuffers.ByteBuffer} bb
* @param {MNN.Blob=} obj
* @returns {MNN.Blob}
*/
MNN.Blob.getRootAsBlob = function(bb, obj) {
return (obj || new MNN.Blob).__init(bb.readInt32(bb.position()) + bb.position(), bb);
};
/**
* @param {flatbuffers.ByteBuffer} bb
* @param {MNN.Blob=} obj
* @returns {MNN.Blob}
*/
MNN.Blob.getSizePrefixedRootAsBlob = function(bb, obj) {
return (obj || new MNN.Blob).__init(bb.readInt32(bb.position()) + bb.position(), bb);
};
具体关于 FlatBuffer 的细节，可以阅读我之前的文章，不再赘述。
避免冗余解析流程
上文提到 根据 OpParameter 来获取 main 中的数据，然后依次填入 parameter / argument 是比较耗费精力的步骤。我们所有的 OpParameter 类型有 74种（还在不断更新）
MNN.OpParameter = {
NONE: 0,
QuantizedAdd: 1,
ArgMax: 2,
AsString: 3,
Axis: 4,
BatchNorm: 5,
BinaryOp: 6,
Blob: 7,
CastParam: 8,
Convolution2D: 9,
Crop: 10,
CropAndResize: 11,
Dequantize: 12,
DetectionOutput: 13,
Eltwise: 14,
ExpandDims: 15,
Fill: 16,
Flatten: 17,
Gather: 18,
GatherV2: 19,
InnerProduct: 20,
Input: 21,
Interp: 22,
LRN: 23,
LSTM: 24,
MatMul: 25,
NonMaxSuppressionV2: 26,
Normalize: 27,
PackParam: 28,
Permute: 29,
Plugin: 30,
Pool: 31,
PRelu: 32,
PriorBox: 33,
Proposal: 34,
QuantizedAvgPool: 35,
QuantizedBiasAdd: 36,
QuantizedConcat: 37,
QuantizedLogistic: 38,
QuantizedMatMul: 39,
QuantizedMaxPool: 40,
QuantizedRelu: 41,
QuantizedRelu6: 42,
QuantizedReshape: 43,
QuantizedSoftmax: 44,
QuantizeMaxMin: 45,
QuantizeV2: 46,
Range: 47,
Rank: 48,
ReduceJoin: 49,
ReductionParam: 50,
Relu: 51,
Relu6: 52,
RequantizationRange: 53,
Requantize: 54,
Reshape: 55,
Resize: 56,
RoiPooling: 57,
Scale: 58,
Selu: 59,
Size: 60,
Slice: 61,
SliceTf: 62,
SpaceBatch: 63,
SqueezeParam: 64,
StridedSliceParam: 65,
TensorConvertInfo: 66,
TfQuantizedConv2D: 67,
TopKV2: 68,
Transpose: 69,
UnaryOp: 70,
MomentsParam: 71,
RNNParam: 72,
BatchMatMulParam: 73,
QuantizedFloatParam: 74
};
以 Convolution2D 举例，它又有几个对应的参数：weight ，bias ，quanParameter ，symmetricQuan ，padX ，padY ，kernelX ，kernelY 等等，需要解析。
一开始我采用了人肉的解析方式，代码就成了 if else 加上一大堆解析代码：
mnn_private.Convolution2DAttrBuilder = class {
constructor() {}
buildAttributes(metadata, parameter) {
//var common = parameter.common();
var attributes = [];
var common = parameter.common();
attributes.push(new mnn.Attribute(metadata, "padX", common.padX(), true));
attributes.push(new mnn.Attribute(metadata, "padY", common.padY(), true));
attributes.push(new mnn.Attribute(metadata, "kernelX", common.kernelX(), true));
attributes.push(new mnn.Attribute(metadata, "kernelY", common.kernelY(), true));
attributes.push(new mnn.Attribute(metadata, "strideX", common.strideX(), true));
attributes.push(new mnn.Attribute(metadata, "strideY", common.strideY(), true));
attributes.push(new mnn.Attribute(metadata, "dilateX", common.dilateX(), true));
attributes.push(new mnn.Attribute(metadata, "dilateY", common.dilateY(), true));
attributes.push(new mnn.Attribute(metadata, "padMode", mnn.schema.PadModeName[common.dilateY()], true));
attributes.push(new mnn.Attribute(metadata, "group", common.group(), true));
attributes.push(new mnn.Attribute(metadata, "outputCount", common.outputCount(), true));
attributes.push(new mnn.Attribute(metadata, "inputCount", common.inputCount(), true));
attributes.push(new mnn.Attribute(metadata, "relu", common.relu(), true));
attributes.push(new mnn.Attribute(metadata, "relu6", common.relu6(), true));
//var quanParameter = parameter.quanParameter();
var weights = [];
for (var w = 0; w < parameter.weightLength(); w++) {
weights.push(parameter.weight(w));
}
attributes.push(new mnn.Attribute(metadata, "weights", weights, true));
var bias = [];
for (var b = 0; b < parameter.biasLength(); b++) {
bias.push(parameter.bias(b));
}
attributes.push(new mnn.Attribute(metadata, "bias", bias, true));
return attributes;
}
get hasMain() {
return true;
}
这样的代码如果写完74个 OpParameter ，可维护性和后续的扩展也不够。
我们要巧用 JavaScript的 Reflect 能力以及属性等于与字符串值属性的特性
_buildAttributes(metadata, op, net, args) {
var opParameter = op.mainType();
var opParameterName = mnn.schema.OpParameterName[opParameter];
// 获取对应的类型
var mainConstructor = mnn.schema[opParameterName];
var opParameterObject = null;
if (typeof mainConstructor === 'function') {
var mainTemplate = Reflect.construct(mainConstructor, []);
opParameterObject = op.main(mainTemplate);
}
this._recursivelyBuildAttributes(metadata, net, opParameterObject, this._attributes);
}
_recursivelyBuildAttributes(metadata, net, opParameterObject, attributeHolders) {
if (!opParameterObject) return;
var attributeName;
var attributeNames = [];
var attributeNamesMap = {};
for (attributeName of Object.keys(Object.getPrototypeOf(opParameterObject))) {
if (attributeName != '__init') {
attributeNames.push(attributeName);
}
attributeNamesMap[attributeName] = true;
}
var attributeArrayNamesMap = {};
for (attributeName of Object.keys(attributeNamesMap)) {
if (attributeNamesMap[attributeName + 'Length']) { attributeArrayNamesMap[attributeName] = true;
attributeNames = attributeNames.filter((item) => item != (attributeName + 'Array') && item != (attributeName + 'Length'));
}
}
for (attributeName of attributeNames) {
if (opParameterObject[attributeName] && typeof opParameterObject[attributeName] == 'function') {
var value = null;
if (attributeArrayNamesMap[attributeName]) {
var array = [];
var length = opParameterObject[attributeName + 'Length']();
//var a = opParameterObject[attributeName + 'Array']();
for (var l = 0; l < length; l++) {
array.push(opParameterObject[attributeName + 'Length'](l));
}
value = array;
}
else {
value = opParameterObject[attributeName]();
if (typeof value === 'object') {
this._recursivelyBuildAttributes(metadata, net, value, attributeHolders);
value = null;
}
}
if (value) {
var attribute = new mnn.Attribute(metadata, attributeName, value);
attributeHolders.push(attribute);
}
}
}
}
区区50多行代码就可以完成所有 OpParamater 及其对应的属性解析。
登录
