Add jit transformer resize

Saurav Agarwal · Saurav Agarwal · commit 7cfd12b5cbeb · 2023-05-24T18:50:04.000-04:00
diff --git a/python/tests/data_load.py b/python/tests/data_load.py
@@ -0,0 +1,16 @@
+import sys
+import torch
+import pyCoverageControlTorch as cct
+
+if __name__ == "__main__":
+    # Get filename from argument
+    filename = sys.argv[1]
+    # Load tensor from file
+    tensor_model = torch.jit.load(filename)
+    tensor = list(tensor_model.parameters())[0]
+    # Print tensor
+    print(tensor.dtype)
+    print(tensor.shape)
+    print(tensor.sum())
+
+
diff --git a/python/ts_jit/TorchVisionResize_32.pt b/python/ts_jit/TorchVisionResize_32.pt
diff --git a/python/ts_jit/torchvision_resize.py b/python/ts_jit/torchvision_resize.py
@@ -0,0 +1,21 @@
+import os
+import sys
+import torch
+import torchvision
+import torchvision.transforms as T
+
+class Resizer(torch.nn.Module):
+    def __init__(self, size):
+        super().__init__()
+        self.size = size
+        self.T = T.Resize(size)
+
+    def forward(self, img):
+        return self.T(img)
+
+if __name__ == "__main__":
+    # Get size from args
+    map_size = int(sys.argv[1])
+    file_name = str(sys.argv[2])
+    scripted_resizer = torch.jit.script(Resizer([map_size,]))
+    scripted_resizer.save(file_name)
