Remove gradient calculation during inference (#258)

* Remove force_grad in inference * Upgrade required PyTorch version * Update CHANGELOG.md * Update CHANGELOG.md * Fix typo in torch version * Specify correct Pytorch version change --------- Co-authored-by: Wout Bittremieux <[email protected]>
Noble-Lab · Oct 24, 2023 · 7721962 · 7721962
1 parent 3a18fed
commit 7721962
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 29 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,14 +16,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Changed
 
 - The CLI has been overhauled to use subcommands.
-- Upgraded to Lightning >=2.0
+- Upgraded to Lightning >=2.0.
 - Checkpointing is configured to save the top-k models instead of all.
 - Log steps rather than epochs as units of progress during training.
 - Validation performance metrics are logged (and added to tensorboard) at the validation epoch, and training loss is logged at the end of training epoch, i.e. training and validation metrics are logged asynchronously.
 - Irrelevant warning messages on the console output and in the log file are no longer shown.
 - Nicely format logged warnings.
 - `every_n_train_steps` has been renamed to `val_check_interval` in accordance to the corresponding Pytorch Lightning parameter.
 - Training batches are randomly shuffled.
+- Upgraded to Torch >=2.1.
 
 ### Removed
 
@@ -36,7 +37,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Correctly refer to input peak files by their full file path.
 - Specifying custom residues to retrain Casanovo is now possible.
 - Upgrade to depthcharge v0.2.3 to fix sinusoidal encoding and for the `PeptideTransformerDecoder` hotfix.
-- Enable gradients during prediction and validation to avoid NaNs from occuring as a temporary workaround until a new Pytorch version is available.
 - Correctly report amino acid precision and recall during validation.
 
 ## [3.3.0] - 2023-04-04

diff --git a/casanovo/denovo/model.py b/casanovo/denovo/model.py
@@ -751,9 +751,7 @@ def validation_step(
             The loss of the validation step.
         """
         # Record the loss.
-        # FIXME: Temporary workaround to avoid the NaN bug.
-        with torch.set_grad_enabled(True):
-            loss = self.training_step(batch, mode="valid")
+        loss = self.training_step(batch, mode="valid")
         if not self.calculate_precision:
             return loss
 
@@ -804,30 +802,28 @@ def predict_step(
             and amino acid-level confidence scores.
         """
         predictions = []
-        # FIXME: Temporary workaround to avoid the NaN bug.
-        with torch.set_grad_enabled(True):
-            for (
-                precursor_charge,
-                precursor_mz,
-                spectrum_i,
-                spectrum_preds,
-            ) in zip(
-                batch[1][:, 1].cpu().detach().numpy(),
-                batch[1][:, 2].cpu().detach().numpy(),
-                batch[2],
-                self.forward(batch[0], batch[1]),
-            ):
-                for peptide_score, aa_scores, peptide in spectrum_preds:
-                    predictions.append(
-                        (
-                            spectrum_i,
-                            precursor_charge,
-                            precursor_mz,
-                            peptide,
-                            peptide_score,
-                            aa_scores,
-                        )
+        for (
+            precursor_charge,
+            precursor_mz,
+            spectrum_i,
+            spectrum_preds,
+        ) in zip(
+            batch[1][:, 1].cpu().detach().numpy(),
+            batch[1][:, 2].cpu().detach().numpy(),
+            batch[2],
+            self.forward(batch[0], batch[1]),
+        ):
+            for peptide_score, aa_scores, peptide in spectrum_preds:
+                predictions.append(
+                    (
+                        spectrum_i,
+                        precursor_charge,
+                        precursor_mz,
+                        peptide,
+                        peptide_score,
+                        aa_scores,
                     )
+                )
 
         return predictions
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -34,7 +34,7 @@ dependencies = [
     "scikit-learn",
     "spectrum_utils",
     "tensorboard",
-    "torch>=2.0",
+    "torch>=2.1",
     "tqdm",
 ]
 dynamic = ["version"]