Force gradient calculation during inference (#231)

* Add lines to force gradient calculation. During inference only. * Add comments in the code to document the temporary workaround * Update changelog --------- Co-authored-by: Varun Ananth <[email protected]> Co-authored-by: Wout Bittremieux <[email protected]> Co-authored-by: Wout Bittremieux <[email protected]>
Noble-Lab · Aug 18, 2023 · d9396aa · d9396aa
1 parent 51221f9
commit d9396aa
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 22 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,17 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased]
 
+### Changed
+
+- The CLI has been overhauled to use subcommands.
+- Upgraded to Lightning >=2.0
+- Checkpointing is now configured to save the top-k models instead of all.
+
+### Fixed
+
+- Casanovo now runs on CPU and can passes all tests.
+- Enable gradients during prediction and validation to avoid NaNs from occuring as a temporary workaround until a new Pytorch version is available.
+
 ### Added
 
 -  Checkpoints now include model parameters, allowing for mismatches with the provided configuration file.

diff --git a/casanovo/denovo/model.py b/casanovo/denovo/model.py
@@ -758,7 +758,9 @@ def validation_step(
             The loss of the validation step.
         """
         # Record the loss.
-        loss = self.training_step(batch, mode="valid")
+        # FIXME: Temporary workaround to avoid the NaN bug.
+        with torch.set_grad_enabled(True):
+            loss = self.training_step(batch, mode="valid")
         if not self.calculate_precision:
             return loss
 
@@ -809,28 +811,30 @@ def predict_step(
             and amino acid-level confidence scores.
         """
         predictions = []
-        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,
+        # 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,
+                        )
                     )
-                )
 
         return predictions