Remove jit from metadata.

pytorch · Jul 25, 2023 · 98baa08 · 98baa08
1 parent 1d2e3bf
commit 98baa08
Show file tree

Hide file tree

Showing 56 changed files with 15 additions and 109 deletions.
diff --git a/torchbenchmark/models/Background_Matting/__init__.py b/torchbenchmark/models/Background_Matting/__init__.py
@@ -100,21 +100,15 @@ def __init__(self, test, device, batch_size=None, extra_args=[]):
         self.log_writer = SummaryWriter(datadir)
         self.model_dir = datadir
 
-        self._maybe_trace()
-
-    def _maybe_trace(self):
+    def jit_callback(self):
         for data in self.train_data:
             bg, image, seg, multi_fr = data['bg'], data['image'], data['seg'], data['multi_fr']
             bg, image, seg, multi_fr = Variable(bg.to(self.device)), Variable(
                 image.to(self.device)), Variable(seg.to(self.device)), Variable(multi_fr.to(self.device))
-            if self.jit:
-                self.netB = torch.jit.trace(
-                    self.netB, (image, bg, seg, multi_fr))
-                self.netG = torch.jit.trace(
-                    self.netG, (image, bg, seg, multi_fr))
-            else:
-                self.netB(image, bg, seg, multi_fr)
-                self.netG(image, bg, seg, multi_fr)
+            self.netB = torch.jit.trace(
+                self.netB, (image, bg, seg, multi_fr))
+            self.netG = torch.jit.trace(
+                self.netG, (image, bg, seg, multi_fr))
             break
 
     def get_module(self):

diff --git a/torchbenchmark/models/LearningToPaint/metadata.yaml b/torchbenchmark/models/LearningToPaint/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: true
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: true
 train_deterministic: false
diff --git a/torchbenchmark/models/basic_gnn_edgecnn/metadata.yaml b/torchbenchmark/models/basic_gnn_edgecnn/metadata.yaml
@@ -3,5 +3,4 @@ eval_deterministic: false
 eval_nograd: true
 train_benchmark: false
 train_deterministic: false
-not_implemented:
-  - jit: true
+
diff --git a/torchbenchmark/models/basic_gnn_gcn/metadata.yaml b/torchbenchmark/models/basic_gnn_gcn/metadata.yaml
@@ -3,5 +3,3 @@ eval_deterministic: false
 eval_nograd: true
 train_benchmark: false
 train_deterministic: false
-not_implemented:
-  - jit: true
diff --git a/torchbenchmark/models/basic_gnn_gin/metadata.yaml b/torchbenchmark/models/basic_gnn_gin/metadata.yaml
@@ -2,6 +2,4 @@ eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
 train_benchmark: false
-train_deterministic: false
-not_implemented:
-  - jit: true
+train_deterministic: false
diff --git a/torchbenchmark/models/basic_gnn_sage/metadata.yaml b/torchbenchmark/models/basic_gnn_sage/metadata.yaml
@@ -3,5 +3,3 @@ eval_deterministic: false
 eval_nograd: true
 train_benchmark: false
 train_deterministic: false
-not_implemented:
-  - jit: true
diff --git a/torchbenchmark/models/dcgan/__init__.py b/torchbenchmark/models/dcgan/__init__.py
@@ -92,8 +92,6 @@ def __init__(self, dcgan):
             # state size. (dcgan.nc) x 64 x 64
         )
 
-        self.jt = None
-        self.jitshape = None
         self.debug_print = False
 
     def forward(self, input):
@@ -131,8 +129,6 @@ def __init__(self, ncgan):
             nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
             nn.Sigmoid()
         )
-        self.jt = None
-        self.jitshape = None
 
     def forward(self, input):
       return self.main(input)

diff --git a/torchbenchmark/models/detectron2_fasterrcnn_r_101_c4/metadata.yaml b/torchbenchmark/models/detectron2_fasterrcnn_r_101_c4/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_fasterrcnn_r_101_dc5/metadata.yaml b/torchbenchmark/models/detectron2_fasterrcnn_r_101_dc5/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_fasterrcnn_r_101_fpn/metadata.yaml b/torchbenchmark/models/detectron2_fasterrcnn_r_101_fpn/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_fasterrcnn_r_50_c4/metadata.yaml b/torchbenchmark/models/detectron2_fasterrcnn_r_50_c4/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_fasterrcnn_r_50_dc5/metadata.yaml b/torchbenchmark/models/detectron2_fasterrcnn_r_50_dc5/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_fasterrcnn_r_50_fpn/metadata.yaml b/torchbenchmark/models/detectron2_fasterrcnn_r_50_fpn/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_fcos_r_50_fpn/metadata.yaml b/torchbenchmark/models/detectron2_fcos_r_50_fpn/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_maskrcnn/metadata.yaml b/torchbenchmark/models/detectron2_maskrcnn/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_maskrcnn_r_101_c4/metadata.yaml b/torchbenchmark/models/detectron2_maskrcnn_r_101_c4/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_maskrcnn_r_101_fpn/metadata.yaml b/torchbenchmark/models/detectron2_maskrcnn_r_101_fpn/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_maskrcnn_r_50_c4/metadata.yaml b/torchbenchmark/models/detectron2_maskrcnn_r_50_c4/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/detectron2_maskrcnn_r_50_fpn/metadata.yaml b/torchbenchmark/models/detectron2_maskrcnn_r_50_fpn/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
 - device: cpu
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/dlrm/metadata.yaml b/torchbenchmark/models/dlrm/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: true
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: true
diff --git a/torchbenchmark/models/doctr_det_predictor/metadata.yaml b/torchbenchmark/models/doctr_det_predictor/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
   - device: cpu
-  - jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/doctr_reco_predictor/metadata.yaml b/torchbenchmark/models/doctr_reco_predictor/metadata.yaml
@@ -6,6 +6,5 @@ eval_deterministic: false
 eval_nograd: true
 not_implemented:
   - device: cpu
-  - jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/drq/metadata.yaml b/torchbenchmark/models/drq/metadata.yaml
@@ -2,7 +2,4 @@ eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
 train_benchmark: false
-train_deterministic: false
-not_implemented:
-  # DrQ model does not support JIT
-  - jit: true
+train_deterministic: false
diff --git a/torchbenchmark/models/fastNLP_Bert/metadata.yaml b/torchbenchmark/models/fastNLP_Bert/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: true
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: true
diff --git a/torchbenchmark/models/hf_Albert/metadata.yaml b/torchbenchmark/models/hf_Albert/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_Bart/metadata.yaml b/torchbenchmark/models/hf_Bart/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_Bert/metadata.yaml b/torchbenchmark/models/hf_Bert/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_Bert_large/metadata.yaml b/torchbenchmark/models/hf_Bert_large/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_BigBird/metadata.yaml b/torchbenchmark/models/hf_BigBird/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_DistilBert/metadata.yaml b/torchbenchmark/models/hf_DistilBert/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_GPT2/metadata.yaml b/torchbenchmark/models/hf_GPT2/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_GPT2_large/metadata.yaml b/torchbenchmark/models/hf_GPT2_large/metadata.yaml
@@ -4,8 +4,6 @@ eval_nograd: true
 train_benchmark: false
 train_deterministic: false
 not_implemented:
-  # hf_GPT2 model doesn't support JIT
-  - jit: true
   # OOMs on torchbench CI
   - device: cuda
   # CPU OOM on torchbench CI

diff --git a/torchbenchmark/models/hf_Longformer/metadata.yaml b/torchbenchmark/models/hf_Longformer/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_Reformer/metadata.yaml b/torchbenchmark/models/hf_Reformer/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_T5_generate/metadata.yaml b/torchbenchmark/models/hf_T5_generate/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/hf_T5_large/metadata.yaml b/torchbenchmark/models/hf_T5_large/metadata.yaml
@@ -4,7 +4,5 @@ eval_nograd: true
 train_benchmark: false
 train_deterministic: false
 not_implemented:
-  # hf_T5 model doesn't support JIT
-  - jit: true
   # disable train test because of CI infra capacity issue
   - test: train
diff --git a/torchbenchmark/models/lennard_jones/metadata.yaml b/torchbenchmark/models/lennard_jones/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/llama_v2_7b_16h/metadata.yaml b/torchbenchmark/models/llama_v2_7b_16h/metadata.yaml
@@ -5,7 +5,6 @@ eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
 not_implemented:
-- jit: true
 - device: cpu
 - device: cuda
   test: train

diff --git a/torchbenchmark/models/maml/metadata.yaml b/torchbenchmark/models/maml/metadata.yaml
@@ -2,6 +2,4 @@ eval_benchmark: false
 eval_deterministic: true
 eval_nograd: true
 train_benchmark: false
-train_deterministic: true
-not_implemented:
-  - jit: true
+train_deterministic: true
diff --git a/torchbenchmark/models/maml_omniglot/metadata.yaml b/torchbenchmark/models/maml_omniglot/metadata.yaml
@@ -2,6 +2,4 @@ eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
 train_benchmark: false
-train_deterministic: false
-not_implemented:
-  - jit: true
+train_deterministic: false
diff --git a/torchbenchmark/models/mobilenet_v2_quantized_qat/__init__.py b/torchbenchmark/models/mobilenet_v2_quantized_qat/__init__.py
@@ -19,8 +19,6 @@ class Model(BenchmarkModel):
     def __init__(self, test, device, batch_size=None, extra_args=[]):
         if test == "eval" and device != "cpu":
             raise NotImplementedError("The eval test only supports CPU.")
-        if jit and test == "train":
-            raise NotImplementedError("torchscript operations should only be applied after quantization operations")
         super().__init__(test=test, device=device, batch_size=batch_size, extra_args=extra_args)
 
         self.model = models.mobilenet_v2().to(self.device)

diff --git a/torchbenchmark/models/moco/__init__.py b/torchbenchmark/models/moco/__init__.py
@@ -88,7 +88,6 @@ def collate_train_fn(data):
     def get_module(self):
         """ Recommended
         Returns model, example_inputs
-        model should be torchscript model if self.jit is True.
         Both model and example_inputs should be on self.device properly.
         `model(*example_inputs)` should execute one step of model forward.
         """

diff --git a/torchbenchmark/models/moco/metadata.yaml b/torchbenchmark/models/moco/metadata.yaml
@@ -5,7 +5,6 @@ eval_benchmark: true
 eval_deterministic: false
 eval_nograd: true
 not_implemented:
-- jit: true
 - device: cpu
 train_benchmark: true
 train_deterministic: false
diff --git a/torchbenchmark/models/nanogpt_generate/metadata.yaml b/torchbenchmark/models/nanogpt_generate/metadata.yaml
@@ -1,7 +1,5 @@
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/nvidia_deeprecommender/metadata.yaml b/torchbenchmark/models/nvidia_deeprecommender/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/opacus_cifar10/metadata.yaml b/torchbenchmark/models/opacus_cifar10/metadata.yaml
@@ -4,7 +4,5 @@ devices:
 eval_benchmark: false
 eval_deterministic: false
 eval_nograd: true
-not_implemented:
-- jit: true
 train_benchmark: false
 train_deterministic: false
diff --git a/torchbenchmark/models/pytorch_struct/metadata.yaml b/torchbenchmark/models/pytorch_struct/metadata.yaml
@@ -2,6 +2,4 @@ eval_benchmark: true
 eval_deterministic: true
 eval_nograd: true
 train_benchmark: false
-train_deterministic: false
-not_implemented:
-  - jit: true
+train_deterministic: false