Merge pull request #226 from JingyuanZhang/master

feat(core): update model vars structure for high performance

Author: yueshuangyan

packages/paddlejs-converter/convertModel.py

@@ -33,7 +33,7 @@
 # paddlepaddle运行程序实例
 program = None
 # 存放模型结构
-modelInfo = {"vars": [], "ops": [], "chunkNum": 0}
+modelInfo = {"vars": {}, "ops": [], "chunkNum": 0, "dataLayout": "nhwc"}
 # 存放参数数值（未排序）
 paramValuesDict = {}
 
@@ -236,7 +236,7 @@ def organizeModelVariableInfo(result):
     # 将var信息按照顺序，添加到model info的vars中
     for key, value in varInfoOrderDict.items():
         value["name"] = key
-        modelInfo["vars"].append(value)
+        modelInfo["vars"][key] = value
     print("Organizing model variables info successfully.")
 
 def organizeModelOpInfo():

packages/paddlejs-core/src/commons/utils.ts

@@ -28,8 +28,8 @@ export function getGlobalInterface(): any {
 /**
  * getOrMakeGlobalProperty
  * @param {string} key property key
- * @param { any } value property value
- * @returns { any } global property
+ * @param {any} value property value
+ * @returns {any} global property
  */
 export function getOrMakeGlobalProperty(key: string, value?: Object | String | Number | Boolean)
     : any {
@@ -39,4 +39,63 @@ export function getOrMakeGlobalProperty(key: string, value?: Object | String | N
     }
     globalInterface[key] = value;
     return globalInterface[key];
+}
+
+
+/**
+ * find target var by key
+ * @param {Object | Array} vars model vars
+ * @param {string} key var name
+ * @returns {Object} var
+ */
+export function findVarByKey(vars, key) {
+    if (Array.isArray(vars)) {
+        return vars.find(item => item.name === key);
+    }
+    return vars[key];
+}
+
+
+/**
+ * add var to vars
+ * @param {Object | Array} vars model vars
+ * @param {Object} item var
+ */
+export function AddItemToVars(vars, item) {
+    const isVarsArray = Array.isArray(vars);
+    const isItemArray = Array.isArray(item);
+    if (isItemArray) {
+        if (isVarsArray) {
+            vars.splice(vars.length - 1, 0, ...item);
+        }
+        else {
+            item.forEach(varItem => {
+                vars[varItem.name] = varItem;
+            });
+        }
+        return;
+    }
+    if (isVarsArray) {
+        vars.push(item);
+        return;
+    }
+    vars[item.name] = item;
+}
+
+/**
+ * traverse vars and deal var item
+ * @param {Object | Array} vars model vars
+ * @param {Function} callback deal var item
+ */
+export function traverseVars(vars, callback) {
+    const isVarsArray = Array.isArray(vars);
+    if (isVarsArray) {
+        vars.forEach(item => {
+            callback(item);
+        });
+        return;
+    }
+    Object.keys(vars).forEach(key => {
+        callback(vars[key]);
+    });
 }

packages/paddlejs-core/src/index.ts

@@ -4,6 +4,7 @@ import { registerBackend, registerOp } from './globals';
 import Env from './env';
 import * as interfaces from './commons/interface';
 import Transformer from './transform/transformer';
+import * as coreUtils from './commons/utils';
 
 export {
     Runner,
@@ -12,5 +13,6 @@ export {
     PaddlejsBackend,
     interfaces,
     Transformer,
-    Env as env
+    Env as env,
+    coreUtils
 };

packages/paddlejs-core/src/loader.ts

@@ -3,7 +3,8 @@
  */
 
 import env from './env';
-import { Model, ModelVar } from './commons/interface';
+import { Model } from './commons/interface';
+import { traverseVars } from './commons/utils';
 
 interface UrlConf {
     dir: string;
@@ -124,12 +125,10 @@ export default class ModelLoader {
         });
     }
 
-    traverse(arr: ModelVar[], allChunksData: Float32Array) {
+    traverse(vars, allChunksData: Float32Array) {
         let marker = 0; // 读到哪个位置了
         let len; // 当前op长度
-        arr.filter(item => {
-            return item.name;
-        }).forEach(item => {
+        traverseVars(vars, item => {
             len = item.shape.reduce((a, b) => a * b); // 长度为shape的乘积
             // 为了减少模型体积，模型转换工具不会导出非persistable的数据，这里只需要读取persistable的数据
             if (item.persistable) {

packages/paddlejs-core/src/opFactory/opDataBuilder.ts

@@ -3,6 +3,7 @@ import { GLOBALS } from '../globals';
 import Tensor from './tensor';
 import opBehaviors from './opBehaviors';
 import * as Utils from './utils';
+import { findVarByKey } from '../commons/utils';
 
 // model的名字和paddleJS的tensor名字mapping
 
@@ -67,12 +68,12 @@ export default class OpData {
     constructTensorData() {
         Object.keys(this.output).forEach(key => {
             this.output[key].forEach((name: string, index: number) => {
-                this.output[key][index] = this.getTensorVar(name)[0];
+                this.output[key][index] = this.getTensorVar(name);
             });
         });
 
         Object.keys(this.input).forEach(key => {
-            this.input[key] = this.getTensorVar(this.input[key][0]);
+            this.input[key] = [this.getTensorVar(this.input[key][0])];
         });
 
         for (const key in this.output) {
@@ -142,10 +143,11 @@ export default class OpData {
     }
 
     getTensorVar(name: string) {
-        const data = this.vars.filter(item => item.name === name || item.name === name.replace(/_packed$/, ''));
-        if (data.length > 0 && name.endsWith('_packed')) {
-            const packedData = Utils.packOpData(data[0], name);
-            return [packedData];
+        const varName = name.replace(/_packed$/, '');
+        const data = findVarByKey(this.vars, varName);
+        if (data && name.endsWith('_packed')) {
+            const packedData = Utils.packOpData(data, name);
+            return packedData;
         }
         return data;
     }

packages/paddlejs-core/src/runner.ts

@@ -4,7 +4,7 @@ import { Model, ModelConfig, InputFeed, ModelVar, GraphType } from './commons/in
 import OpData from './opFactory/opDataBuilder';
 import Tensor from './opFactory/tensor';
 import { GLOBALS } from './globals';
-import { getGlobalInterface } from './commons/utils';
+import { getGlobalInterface, findVarByKey, AddItemToVars } from './commons/utils';
 import MediaProcessor from './mediaProcessor';
 import env from './env';
 
@@ -174,7 +174,7 @@ export default class Runner {
             if (feedOpInputs.length > 1) {
                 // 多输入
                 preheatFeedData = feedOpInputs.map(inputName => {
-                    const feedInfo = vars.find(item => item.name === inputName);
+                    const feedInfo = findVarByKey(vars, inputName);
                     const shape = feedInfo.shape;
                     const [w, h, c = 3, n = 1] = shape.reverse();
 
@@ -184,7 +184,7 @@ export default class Runner {
             }
         }
         else {
-            preheatFeedData = vars.find(item => item.name === 'image');
+            preheatFeedData = findVarByKey(vars, 'image');
             if (preheatFeedData) {
                 preheatFeedData.data = new Float32Array(fc * fh * fw).fill(1.0);
                 return;
@@ -197,7 +197,7 @@ export default class Runner {
             };
         }
 
-        vars.push(preheatFeedData);
+        AddItemToVars(vars, preheatFeedData);
     }
 
     updateFeedData(feed) {
@@ -266,9 +266,7 @@ export default class Runner {
 
     async read() {
         const fetchOp = this.graphGenerator.getFetchExecutor();
-        const fetchVar = this.model.vars.find(
-            item => item.name === fetchOp.inputs.X[0]
-        ) as ModelVar;
+        const fetchVar = findVarByKey(this.model.vars, fetchOp.inputs.X[0]) as ModelVar;
         const fetchInfo = {
             name: fetchVar.name,
             shape: fetchOp.attrs['origin_shape'] || fetchVar.shape

packages/paddlejs-core/src/transform/feedProcess.ts

@@ -3,6 +3,7 @@
  */
 
 import { ModelOp } from '../commons/interface';
+import { findVarByKey, AddItemToVars } from '../commons/utils';
 import env from '../env';
 import Transformer from './transformer';
 
@@ -25,7 +26,7 @@ export default class WebglFeedProcess extends Transformer {
         } = modelConfig;
 
         // make img_pre_processed var
-        const imgVar = vars.find(item => item.name === 'image');
+        const imgVar = findVarByKey(vars, 'image');
         const [, , h, w] = imgVar.shape;
         imgVar.shape = [1, 1, h, w];
         const processImgVar = Object.assign({}, imgVar);
@@ -40,8 +41,7 @@ export default class WebglFeedProcess extends Transformer {
         originImgVar.persistable = false;
         delete originImgVar.data;
 
-        vars.push(originImgVar);
-        vars.push(processImgVar);
+        AddItemToVars(vars, [originImgVar, processImgVar]);
 
         // change recieve_img op input
         const imageOriginInputOp = ops.find(item => {

packages/paddlejs-core/src/transform/nhwc2nchw.ts

@@ -4,6 +4,7 @@
 
 import env from '../env';
 import { ModelOp } from '../commons/interface';
+import { findVarByKey, AddItemToVars } from '../commons/utils';
 import { formatShape } from '../opFactory/utils';
 import Transformer from './transformer';
 
@@ -21,7 +22,7 @@ export default class nhwc2nchw extends Transformer {
         const [ops, vars] = args;
         const fetchOp = ops.find(item => item.type === 'fetch');
         const [inputName] = fetchOp.inputs.X;
-        const fetchInputVar = vars.find(item => item.name === inputName);
+        const fetchInputVar = findVarByKey(vars, inputName);
         const [n, c, h, w] = formatShape(fetchInputVar.shape);
 
         // transform data from nhwc to nchw
@@ -47,6 +48,6 @@ export default class nhwc2nchw extends Transformer {
 
         fetchOp.inputs.X = [FINAL_NCHW_OP_NAME];
         ops.push(...[nchwOp]);
-        vars.push(...[nchwVar]);
+        AddItemToVars(vars, nchwVar);
     }
 }

packages/paddlejs-core/src/transform/packOutOp.ts

@@ -4,6 +4,7 @@
 
 import env from '../env';
 import { ModelOp } from '../commons/interface';
+import { findVarByKey, AddItemToVars } from '../commons/utils';
 import { formatShape } from '../opFactory/utils';
 import Transformer from './transformer';
 
@@ -22,7 +23,7 @@ export default class PackOut extends Transformer {
         const [ops, vars] = args;
         const fetchOp = ops.find(item => item.type === 'fetch');
         const [inputName] = fetchOp.inputs.X;
-        const fetchInputVar = vars.find(item => item.name === inputName);
+        const fetchInputVar = findVarByKey(vars, inputName);
         const [n, c, h, w] = formatShape(fetchInputVar.shape);
 
         // transform data from nhwc to nchw
@@ -68,7 +69,7 @@ export default class PackOut extends Transformer {
         fetchOp.inputs.X = [FINAL_PACK_OP_NAME];
         fetchOp.attrs['origin_shape'] = [n, c, h, w];
         ops.push(...[nchwOp, packOutOp]);
-        vars.push(...[nchwVar, packOutVar]);
+        AddItemToVars(vars, [nchwVar, packOutVar]);
     }
 }
 

packages/paddlejs-core/src/transform/splitOp.ts

@@ -2,11 +2,13 @@
  * @file graph transformer
  */
 
+import { findVarByKey, AddItemToVars } from '../commons/utils';
 import Transformer from './transformer';
 
+
 function getTensorShapeFromVals(name, vars) {
-    const result = vars.filter(item => item.name === name);
-    return result.length ? result[0].shape : [];
+    const result = findVarByKey(vars, name);
+    return result ? result.shape : [];
 }
 
 function buildOutputVarInfo(inputs, outputShape, axis, vars) {
@@ -83,7 +85,8 @@ export default class SplitOp extends Transformer {
             // change outputname of next op
             opList[opLen - 1].outputs.Out = [outputName];
             ops.splice(index, 1, ...opList);
-            vars.splice(vars.length - 1, 0, ...varList);
+
+            AddItemToVars(vars, varList);
         }
     }
 }