Remove all torch.no_grad contexts in eval.

pytorch · May 3, 2024 · 3450ca2 · 3450ca2
1 parent 3a26c36
commit 3450ca2
Show file tree

Hide file tree

Showing 13 changed files with 39 additions and 52 deletions.
diff --git a/torchbenchmark/models/DALLE2_pytorch/__init__.py b/torchbenchmark/models/DALLE2_pytorch/__init__.py
@@ -84,8 +84,7 @@ def set_module(self, new_model):
 
     def eval(self):
         model, inputs = self.get_module()
-        with torch.no_grad():
-            images = model(*inputs)
+        images = model(*inputs)
         return (images,)
 
     def train(self):

diff --git a/torchbenchmark/models/cm3leon_generate/__init__.py b/torchbenchmark/models/cm3leon_generate/__init__.py
@@ -47,6 +47,5 @@ def train(self):
         return loss.item()
 
     def eval(self):
-        with torch.no_grad():
-            out = self.model(*self.example_inputs)
+        out = self.model(*self.example_inputs)
         return (out,)
diff --git a/torchbenchmark/models/functorch_maml_omniglot/__init__.py b/torchbenchmark/models/functorch_maml_omniglot/__init__.py
@@ -104,6 +104,5 @@ def train(self):
     def eval(self) -> Tuple[torch.Tensor]:
         model, (example_input,) = self.get_module()
         model.eval()
-        with torch.no_grad():
-            out = model(example_input)
+        out = model(example_input)
         return (out, )
diff --git a/torchbenchmark/models/hf_Whisper/__init__.py b/torchbenchmark/models/hf_Whisper/__init__.py
@@ -24,8 +24,7 @@ def get_module(self):
 
     def eval(self):
         self.model.eval()
-        with torch.no_grad():
-            self.model(self.example_inputs["input_ids"])
+        self.model(self.example_inputs["input_ids"])
 
     def enable_fp16(self):
         self.model.half()

diff --git a/torchbenchmark/models/lennard_jones/__init__.py b/torchbenchmark/models/lennard_jones/__init__.py
@@ -100,6 +100,5 @@ def train(self):
     def eval(self) -> Tuple[torch.Tensor]:
         model = self.model
         model.eval()
-        with torch.no_grad():
-            out = make_prediction(model, self.drs)
+        out = make_prediction(model, self.drs)
         return out
diff --git a/torchbenchmark/models/maml_omniglot/__init__.py b/torchbenchmark/models/maml_omniglot/__init__.py
@@ -114,6 +114,5 @@ def train(self):
     def eval(self) -> Tuple[torch.Tensor]:
         model, (example_input,) = self.get_module()
         model.eval()
-        with torch.no_grad():
-            out = model(example_input)
+        out = model(example_input)
         return (out, )
diff --git a/torchbenchmark/models/opacus_cifar10/__init__.py b/torchbenchmark/models/opacus_cifar10/__init__.py
@@ -85,6 +85,5 @@ def eval(self) -> Tuple[torch.Tensor]:
         (images, ) = self.example_inputs
         model.eval()
         targets = self.example_target
-        with torch.no_grad():
-            out = model(images)
+        out = model(images)
         return (out, )
diff --git a/torchbenchmark/models/pyhpc_equation_of_state/__init__.py b/torchbenchmark/models/pyhpc_equation_of_state/__init__.py
@@ -53,6 +53,5 @@ def train(self):
 
     def eval(self) -> Tuple[torch.Tensor]:
         model, example_inputs = self.get_module()
-        with torch.no_grad():
-            out = model(*example_inputs)
+        out = model(*example_inputs)
         return (out, )
diff --git a/torchbenchmark/models/pyhpc_isoneutral_mixing/__init__.py b/torchbenchmark/models/pyhpc_isoneutral_mixing/__init__.py
@@ -146,6 +146,5 @@ def train(self):
 
     def eval(self) -> Tuple[torch.Tensor]:
         model, example_inputs = self.get_module()
-        with torch.no_grad():
-            out = model(*example_inputs)
+        out = model(*example_inputs)
         return out
diff --git a/torchbenchmark/models/pyhpc_turbulent_kinetic_energy/__init__.py b/torchbenchmark/models/pyhpc_turbulent_kinetic_energy/__init__.py
@@ -145,6 +145,5 @@ def train(self):
 
     def eval(self) -> Tuple[torch.Tensor]:
         model, example_inputs = self.get_module()
-        with torch.no_grad():
-            out = model(*example_inputs)
+        out = model(*example_inputs)
         return out
diff --git a/torchbenchmark/models/pytorch_unet/__init__.py b/torchbenchmark/models/pytorch_unet/__init__.py
@@ -90,18 +90,17 @@ def jit_callback(self):
     def eval(self) -> Tuple[torch.Tensor]:
         torch.backends.cudnn.deterministic = True
         self.model.eval()
-        with torch.no_grad():
-            with torch.cuda.amp.autocast(enabled=self.args.amp):
-                mask_pred = self.model(self.example_inputs)
-
-                if self.model.n_classes == 1:
-                    mask_pred = (F.sigmoid(mask_pred) > 0.5).float()
-                else:
-                    mask_pred = (
-                        F.one_hot(mask_pred.argmax(dim=1), self.model.n_classes)
-                        .permute(0, 3, 1, 2)
-                        .float()
-                    )
+        with torch.cuda.amp.autocast(enabled=self.args.amp):
+            mask_pred = self.model(self.example_inputs)
+
+            if self.model.n_classes == 1:
+                mask_pred = (F.sigmoid(mask_pred) > 0.5).float()
+            else:
+                mask_pred = (
+                    F.one_hot(mask_pred.argmax(dim=1), self.model.n_classes)
+                    .permute(0, 3, 1, 2)
+                    .float()
+                )
         return (mask_pred,)
 
     def _get_args(self):

diff --git a/torchbenchmark/models/soft_actor_critic/__init__.py b/torchbenchmark/models/soft_actor_critic/__init__.py
@@ -258,21 +258,20 @@ def train(self):
 
     def eval(self) -> Tuple[torch.Tensor]:
         niter = 1
-        with torch.no_grad():
-            discount = 1.0
-            episode_return_history = []
-            for episode in range(niter):
-                episode_return = 0.0
-                state, _info = self.test_env.reset()
-                done, info = False, {}
-                for step_num in range(self.args.max_episode_steps):
-                    if done:
-                        break
-                    action = self.agent.forward(state)
-                    state, reward, done, info, _unused = self.test_env.step(action)
-                    episode_return += reward * (discount**step_num)
-                episode_return_history.append(episode_return)
-            retval = torch.tensor(episode_return_history)
+        discount = 1.0
+        episode_return_history = []
+        for episode in range(niter):
+            episode_return = 0.0
+            state, _info = self.test_env.reset()
+            done, info = False, {}
+            for step_num in range(self.args.max_episode_steps):
+                if done:
+                    break
+                action = self.agent.forward(state)
+                state, reward, done, info, _unused = self.test_env.step(action)
+                episode_return += reward * (discount**step_num)
+            episode_return_history.append(episode_return)
+        retval = torch.tensor(episode_return_history)
         return (torch.tensor(action),)
 
     def get_optimizer(self):

diff --git a/torchbenchmark/models/timm_efficientdet/__init__.py b/torchbenchmark/models/timm_efficientdet/__init__.py
@@ -181,9 +181,8 @@ def train(self):
                 # self.lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
 
     def eval(self) -> Tuple[torch.Tensor]:
-        with torch.no_grad():
-            for input, target in self.loader:
-                with self.amp_autocast():
-                    output = self.model(input, img_info=target)
-                self.evaluator.add_predictions(output, target)
+        for input, target in self.loader:
+            with self.amp_autocast():
+                output = self.model(input, img_info=target)
+            self.evaluator.add_predictions(output, target)
         return (output, )