PyTorch formats (pytorch/pytorch#31877)

src/pytorch.js

@@ -1554,6 +1554,12 @@ pytorch.Execution = class {
         this._registerFunction('torch.jit._pickle.build_tensorlist', function(data) {
             return data;
         });
+        this._registerFunction('torch.jit._pickle.build_tensor_from_id', function(data) {
+            return data;
+        });
+        this._registerFunction('torch.jit._pickle.restore_type_tag', function(value /*, type_str */) {
+            return value;
+        });
         this._registerFunction('torch.keys', function(dict) {
             return Object.keys(dict);
         });
@@ -2621,33 +2627,18 @@ pytorch.Container.Zip = class {
 
     get format() {
         if (this._format === undefined) {
-            const versionEntry = this._entry('version');
-            const version = versionEntry ? this._utf8Decoder.decode(versionEntry.data).split('\n').shift() : '';
-            if (version === '' || version === '1' || version === '2') {
-                if (this._entry('model.json')) {
-                    if (this._entry('attributes.pkl')) {
-                        this._format = 'TorchScript v1.1';
-                    }
-                    else {
-                        this._format = 'TorchScript v1.0';
-                    }
-                }
-                else if (this._entry('data.pkl')) {
-                    if (this._entry('constants.pkl')) {
-                        if (version === '2') {
-                            this._format = 'TorchScript v1.4';
-                        }
-                        else {
-                            this._format = 'TorchScript v1.3';
-                        }
-                    }
-                    else {
-                        this._format = 'PyTorch v1.4';
-                    }
-                }
-            }
-            else {
-                throw new pytorch.Error("Unsupported PyTorch ZIP version '" + version + "'.");
+            if (this._entry('model.json')) {
+                this._format = this._entry('attributes.pkl') ? 'TorchScript v1.1' : 'TorchScript v1.0';
+            }
+            else if (this._entry('data.pkl')) {
+                const versionEntry = this._entry('version');
+                const versionNumber = versionEntry ? this._utf8Decoder.decode(versionEntry.data).split('\n').shift() : '';
+                const versionTable = { '1': 'v1.3', '2': 'v1.4', '3': 'v1.5' };
+                const version = versionTable[versionNumber];
+                if (!version) {
+                    this._exceptionCallback(new pytorch.Error("Unsupported PyTorch ZIP version '" + versionNumber + "'."));
+                }
+                this._format = (this._entry('constants.pkl') ? 'TorchScript' : 'PyTorch') + ' ' + (version || 'v#' + versionNumber.toString() );
             }
         }
         return this._format;

test/models.json

@@ -3871,21 +3871,21 @@
     "type":   "pytorch",
     "target": "alexnet.zip.pth",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v1.4",
+    "format": "PyTorch v1.5",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "alexnet.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
+    "format": "TorchScript v1.5",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "alexnet_traced.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
+    "format": "TorchScript v1.5",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -3961,15 +3961,13 @@
     "type":   "pytorch",
     "target": "densenet161.zip.pth",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "densenet161.pt",
-    "render": "skip",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
+    "render": "skip",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -3977,7 +3975,6 @@
     "target": "densenet161_traced.pt",
     "render": "skip",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -3996,29 +3993,26 @@
   {
     "type":   "pytorch",
     "target": "inception_v3.pkl.pth",
-    "link":   "https://pytorch.org/docs/stable/torchvision/models.html",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v0.1.10"
+    "format": "PyTorch v0.1.10",
+    "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "inception_v3.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "inception_v3.zip.pth",
-    "link":   "https://pytorch.org/docs/stable/torchvision/models.html",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v1.4"
+    "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "inception_v3_traced.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -4108,14 +4102,12 @@
     "type":   "pytorch",
     "target": "mobilenet_v2.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "mobilenet_v2_traced.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -4205,29 +4197,26 @@
   {
     "type":   "pytorch",
     "target": "resnet18.pkl.pth",
-    "link":   "https://pytorch.org/docs/stable/torchvision/models.html",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v0.1.10"
+    "format": "PyTorch v0.1.10",
+    "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "resnet18.zip.pth",
-    "link":   "https://pytorch.org/docs/stable/torchvision/models.html",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v1.4"
+    "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "resnet18.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "resnet18_traced.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -4268,29 +4257,25 @@
   {
     "type":   "pytorch",
     "target": "resnet101.pkl.pth",
-    "link":   "https://pytorch.org/docs/stable/torchvision/models.html",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v0.1.10"
+    "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "resnet101.zip.pth",
-    "link":   "https://pytorch.org/docs/stable/torchvision/models.html",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v1.4"
+    "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "resnet101.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "resnet101_traced.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -4339,7 +4324,6 @@
     "type":   "pytorch",
     "target": "shufflenet_v2_x1_0.zip.pth",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -4347,7 +4331,6 @@
     "target": "shufflenet_v2_x1_0.pt",
     "render": "skip",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -4367,21 +4350,18 @@
     "type":   "pytorch",
     "target": "squeezenet1_1.zip.pth",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "squeezenet1_1.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "squeezenet1_1_traced.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -4498,7 +4478,6 @@
     "type":   "pytorch",
     "target": "vgg11_bn.zip.pth",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
@@ -4512,21 +4491,18 @@
     "type":   "pytorch",
     "target": "vgg16.zip.pth",
     "script": "./tools/pytorch sync install zoo",
-    "format": "PyTorch v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "vgg16.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {
     "type":   "pytorch",
     "target": "vgg16_traced.pt",
     "script": "./tools/pytorch sync install zoo",
-    "format": "TorchScript v1.4",
     "link":   "https://pytorch.org/docs/stable/torchvision/models.html"
   },
   {