Add TVM test files (#1342)

lutzroeder · Sep 7, 2024 · 9fe7afa · 9fe7afa
1 parent 7bcebe4
commit 9fe7afa
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 29 deletions.
diff --git a/source/armnn.js b/source/armnn.js
@@ -49,7 +49,7 @@ armnn.ModelFactory = class {
                 break;
             }
             default: {
-                throw new armnn.Error(`Unsupported Arm NN '${context.type}'.`);
+                throw new armnn.Error(`Unsupported Arm NN format '${context.type}'.`);
             }
         }
         const metadata = await context.metadata('armnn-metadata.json');

diff --git a/source/mxnet.js b/source/mxnet.js
@@ -8,27 +8,20 @@ mxnet.ModelFactory = class {
     match(context) {
         const identifier = context.identifier;
         const extension = identifier.split('.').pop().toLowerCase();
-        switch (extension) {
-            case 'json': {
-                const obj = context.peek('json');
-                if (obj && obj.nodes && obj.arg_nodes && obj.heads) {
-                    context.type = 'mxnet.json';
-                    context.target = obj;
-                }
-                break;
-            }
-            case 'params': {
-                const stream = context.stream;
-                const signature = [0x12, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
-                if (stream && stream.length > signature.length && stream.peek(signature.length).every((value, index) => value === signature[index])) {
-                    context.type = 'mxnet.params';
-                }
-                break;
-            }
-            default: {
-                break;
+        if (extension === 'json') {
+            const obj = context.peek('json');
+            if (obj && Array.isArray(obj.nodes) && Array.isArray(obj.arg_nodes) && Array.isArray(obj.heads) /* &&
+                !obj.nodes.some((node) => node && node.op === 'tvm_op') */) {
+                context.type = 'mxnet.json';
+                context.target = obj;
+                return;
             }
         }
+        const stream = context.stream;
+        const signature = [0x12, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+        if (stream && stream.length > signature.length && stream.peek(signature.length).every((value, index) => value === signature[index])) {
+            context.type = 'mxnet.params';
+        }
     }
 
     filter(context, type) {
@@ -182,7 +175,7 @@ mxnet.ModelFactory = class {
                     manifest.format = `MXNet${version ? ` v${version}` : ''}`;
                 }
                 if (symbol.nodes && symbol.nodes.some((node) => node && node.op === 'tvm_op')) {
-                    manifest.producer  = 'TVM';
+                    manifest.format  = 'TVM';
                 }
             }
             return new mxnet.Model(metadata, manifest, symbol, parameters);

diff --git a/source/tvm.js b/source/tvm.js
@@ -0,0 +1,50 @@
+
+const tvm = {};
+
+tvm.ModelFactory = class {
+
+    match(context) {
+        const identifier = context.identifier;
+        const extension = identifier.split('.').pop().toLowerCase();
+        if (extension === 'json') {
+            const obj = context.peek('json');
+            if (obj && Array.isArray(obj.nodes) && Array.isArray(obj.arg_nodes) && Array.isArray(obj.heads) &&
+                obj.nodes.every((node) => node && (node.op === 'null' || node.op === 'tvm_op'))) {
+                context.type = 'tvm.json';
+                context.target = obj;
+                return;
+            }
+        }
+        const stream = context.stream;
+        const signature = [0xB7, 0x9C, 0x04, 0x05, 0x4F, 0x8D, 0xE5, 0xF7];
+        if (stream && signature.length <= stream.length && stream.peek(signature.length).every((value, index) => value === signature[index])) {
+            context.type = 'tvm.params';
+        }
+    }
+
+    filter(context, type) {
+        return context.type !== 'tvm.json' || type !== 'tvm.params';
+    }
+
+    async open(context) {
+        switch (context.type) {
+            case 'tvm.json':
+                throw new tvm.Error(`Unsupported TVN model.`);
+            case 'tvm.params':
+                // https://github.com/apache/tvm/blob/main/src/runtime/file_utils.cc#L184
+                throw new tvm.Error(`Invalid file content. File contains TVN NDArray data.`);
+            default:
+                throw new tvm.Error(`Unsupported TVN format '${context.type}'.`);
+        }
+    }
+};
+
+tvm.Error = class extends Error {
+
+    constructor(message) {
+        super(message);
+        this.name = 'Error loading TVM model.';
+    }
+};
+
+export const ModelFactory = tvm.ModelFactory;
diff --git a/source/view.js b/source/view.js
@@ -5671,6 +5671,7 @@ view.ModelFactoryService = class {
         this.register('./hailo', ['.hn', '.har', '.metadata.json']);
         this.register('./nnc', ['.nnc','.tflite']);
         this.register('./safetensors', ['.safetensors', '.safetensors.index.json']);
+        this.register('./tvm', ['.json', '.params']);
         this.register('./modular', ['.maxviz']);
         this.register('./catboost', ['.cbm']);
         this.register('./cambricon', ['.cambricon']);

diff --git a/test/models.json b/test/models.json
@@ -3299,14 +3299,6 @@
     "format":   "MXNet v0.11.0",
     "link":     "https://github.com/awslabs/mxnet-model-server/tree/master/examples/lstm_ptb"
   },
-  {
-    "type":     "mxnet",
-    "target":   "mobilenet-v1-tvm.json",
-    "source":   "https://github.com/lutzroeder/netron/files/2636924/mobilenet-v1-tvm.json.zip[mobilenet-v1-tvm.json]",
-    "format":   "MXNet",
-    "producer": "TVM",
-    "link":     "https://github.com/lutzroeder/netron/issues/199"
-  },
   {
     "type":     "mxnet",
     "target":   "nin.model",
@@ -7951,6 +7943,28 @@
     "format":   "Torch v7",
     "link":     "https://github.com/cpra/fer-cnn-sota"
   },
+  {
+    "type":     "tvm",
+    "target":   "mobilenet-v1-tvm.json",
+    "source":   "https://github.com/lutzroeder/netron/files/2636924/mobilenet-v1-tvm.json.zip[mobilenet-v1-tvm.json]",
+    "format":   "TVM",
+    "link":     "https://github.com/lutzroeder/netron/issues/1342"
+  },
+  {
+    "type":     "tvm",
+    "target":   "squeezenet_opt_level_1.json,squeezenet_opt_level_1.params",
+    "source":   "https://github.com/user-attachments/files/16919287/squeezenet_opt_level_1.zip[squeezenet_opt_level_1.json,squeezenet_opt_level_1.params]",
+    "format":   "TVM",
+    "link":     "https://github.com/lutzroeder/netron/issues/1342"
+  },
+  {
+    "type":     "tvm",
+    "target":   "squeezenet_opt_level_3.params,squeezenet_opt_level_3.json",
+    "source":   "https://github.com/user-attachments/files/16919289/squeezenet_opt_level_3.zip[squeezenet_opt_level_3.params,squeezenet_opt_level_3.json]",
+    "format":   "TVM",
+    "error":    "Invalid file content. File contains TVN NDArray data.",
+    "link":     "https://github.com/lutzroeder/netron/issues/1342"
+  },
   {
     "type":     "uff",
     "target":   "resnet50_tensorflow_1.6.uff",