Merge pull request #231 from JingyuanZhang/master

fix(core): support old model json without dataLayout

Author: yueshuangyan

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

@@ -32,6 +32,7 @@ export interface ModelVar {
 
 export interface Model {
     chunkNum?: number;
+    dataLayout?: string;
     ops: ModelOp[];
     vars: ModelVar[];
     multiOutputs?: ModelVar[]

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

@@ -1,4 +1,7 @@
-import { ModelVar, OpExecutor, OpInputs, OpOutputs, AttrsData, BufferType, OpUniform } from '../commons/interface';
+import {
+    ModelVar, Model, OpExecutor, OpInputs, OpOutputs,
+    AttrsData, BufferType, OpUniform
+} from '../commons/interface';
 import { GLOBALS } from '../globals';
 import Tensor from './tensor';
 import opBehaviors from './opBehaviors';
@@ -21,14 +24,15 @@ export default class OpData {
     outputTensors: Tensor[] = [];
     fShaderParams: object[] = [];
     vars: ModelVar[] = [];
+    dataLayout: string = '';
     iLayer: number = 0;
     program: string[] = [];
     tensorData: ModelVar[] = [];
     isFinalOp: boolean = false;
     modelName: string;
     bufferType: BufferType = BufferType.FrameBuffer;
 
-    constructor(op: OpExecutor, iLayer: number, vars: ModelVar[], isFinalOp: boolean, modelName: string) {
+    constructor(op: OpExecutor, iLayer: number, model: Model, isFinalOp: boolean, modelName: string) {
         const {
             type,
             inputs,
@@ -46,7 +50,8 @@ export default class OpData {
         this.realName = type;
         this.isPackedOp = isPacked;
         this.bufferType = bufferType;
-        this.vars = vars;
+        this.vars = model.vars;
+        this.dataLayout = model.dataLayout || '';
         this.iLayer = iLayer;
         this.isFinalOp = isFinalOp;
         this.input = inputs;
@@ -221,7 +226,8 @@ export default class OpData {
                 interpType: data.interpType || 'NEAREST',
                 isPacked: this.isPackedOp || data.packed || false,
                 binding: index,
-                noLayout: GLOBALS.backendInstance?.noLayout
+                noLayout: GLOBALS.backendInstance?.noLayout,
+                dataLayout: this.dataLayout
             });
             if (tensorName === 'out') {
                 this.outputTensors.push(tensor);

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

@@ -15,6 +15,7 @@ interface TensorParams {
     isPacked?: boolean;
     binding?: number;
     noLayout?: boolean;
+    dataLayout?: string;
 }
 
 export default class Tensor {
@@ -30,6 +31,7 @@ export default class Tensor {
     data: Float32Array | number[] | Uint8Array | null = null;
     persistable: boolean = false;
     interpType: string = 'NEAREST';
+    dataLayout: string = '';
 
     constructor(opts: TensorParams) {
         this.opts = opts;
@@ -41,6 +43,8 @@ export default class Tensor {
         this.interpType = opts.interpType || 'NEAREST';
         // 设置 tensorId
         this.tensorId = opts.type;
+        // set dataLayout
+        this.dataLayout = opts.dataLayout;
         // 保留 model 原生 shape 长度
         this.unformattedShapeLength = opts.shape.length;
         // tensor的形状
@@ -61,7 +65,7 @@ export default class Tensor {
         this.exceedMax = exceedMax;
         // tensor数据
         if (opts.data && opts.data.length) {
-            this.data = new Float32Array(opts.data);
+            this.data = Utils.genTensorData(opts.data, this.dataLayout, shape, this.isPacked);
             opts.data = null;
         }
     }

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

@@ -127,3 +127,46 @@ export function packOpData(opData, packedName) {
     }
     return packedOpData;
 }
+
+
+/**
+ * 将nchw排布数据转为nhwc排布数据
+ * @param {Array} data tensor data
+ * @param {Array} shape nchw
+ * @returns {Array} nhwc data
+ */
+function nchw2nhwc(data: number[] | Float32Array, shape: number[]): number[] | Float32Array {
+    const [N, C, H, W] = shape;
+    const HXW = H * W;
+    const CXHXW = C * H * W;
+    const nhwcData: number[] | Float32Array = [];
+    for (let n = 0; n < N; n++) {
+        for (let h = 0; h < H; h++) {
+            for (let w = 0; w < W; w++) {
+                for (let c = 0; c < C; c++) {
+                    nhwcData.push(data[n * CXHXW + c * HXW + h * W + w]);
+                }
+            }
+        }
+    }
+    return nhwcData;
+}
+
+/**
+ * 生成 tensor data，如果数据排布为 nhwc 则直接返回 Float32Array，否则进行排布变换
+ * @param {Array} data tensor data
+ * @param {string} dataLayout layout
+ * @param {Array} shape nchw
+ * @param {boolean} isPacked
+ * @returns {Float32Array} nhwc data
+ */
+export function genTensorData(data: number[] | Float32Array, dataLayout: string, shape: number[], isPacked: boolean) {
+    if (dataLayout === 'nhwc') {
+        return new Float32Array(data);
+    }
+    const nhwcData: Float32Array | number[] = nchw2nhwc(
+        data,
+        [shape[0], shape[1] * (isPacked ? 4 : 1), shape[2], shape[3]]
+    );
+    return new Float32Array(nhwcData);
+}

packages/paddlejs-core/src/runner.ts

@@ -75,7 +75,7 @@ export default class Runner {
                 const opData = new OpData(
                     op,
                     iLayer,
-                    this.model.vars,
+                    this.model,
                     isFinalOp,
                     this.modelName
                 );