Merge pull request #13 from ArcInstitute/dhruv/step-fix

Add test step for consistency with model test

Author: MingzeDong

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "arc-stack"
-version = "0.1.1"
+version = "0.1.2"
 description = "Stack is a single-cell foundation model that enables in-context learning at inference time."
 readme = "README.md"
 license = { file = "LICENSE" }

src/stack/finetune/lightning.py

@@ -181,6 +181,34 @@ def validation_step(self, batch: Any, batch_idx: int):  # type: ignore[override]
 
         return {'val_loss': total_loss}
 
+    def test_step(self, batch: Any, batch_idx: int):  # type: ignore[override]
+        result = self._forward_pass_with_teacher(batch)
+
+        total_loss = result['loss']
+        recon_loss = result['recon_loss']
+        mmd_loss = result['mmd_loss']
+        sw_loss = result['sw_loss']
+        cls_loss = result['cls_loss']
+        cls_acc = result['cls_acc']
+
+        self.log('test_loss', total_loss, on_epoch=True, prog_bar=True, sync_dist=True)
+        self.log('test/recon_loss', recon_loss, on_epoch=True, sync_dist=True)
+        self.log('test/mmd_loss', mmd_loss, on_epoch=True, sync_dist=True)
+        self.log('test/sw_loss', sw_loss, on_epoch=True, sync_dist=True)
+        self.log('test/cls_loss', cls_loss, on_epoch=True, sync_dist=True)
+        self.log('test/cls_acc', cls_acc, on_epoch=True, sync_dist=True)
+
+        if 'masked_mae' in result:
+            self.log('test/masked_mae', result['masked_mae'], on_epoch=True, sync_dist=True)
+        if 'masked_corr' in result:
+            self.log('test/masked_corr', result['masked_corr'], on_epoch=True, sync_dist=True)
+        if 'sw_predict' in result:
+            self.log('test/sw_predict', result['sw_predict'], on_epoch=True, sync_dist=True)
+        if 'mask_rate' in result:
+            self.log('test/mask_rate', result['mask_rate'], on_epoch=True, sync_dist=True)
+
+        return {'test_loss': total_loss}
+
     def configure_optimizers(self):  # type: ignore[override]
         optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
         config: Dict[str, Any] = {'optimizer': optimizer}

src/stack/training/lightning.py

@@ -55,6 +55,18 @@ def validation_step(self, batch, batch_idx):  # type: ignore[override]
         self.log("val/mask_rate", result["mask_rate"], on_epoch=True, sync_dist=True)
         return {"val_loss": total_loss}
 
+    def test_step(self, batch, batch_idx):  # type: ignore[override]
+        features, _ = batch
+        result = self.model(features, return_loss=True)
+        total_loss = result["loss"]
+        self.log("test_loss", total_loss, on_epoch=True, prog_bar=True, sync_dist=True)
+        self.log("test/recon_loss", result["recon_loss"], on_epoch=True, sync_dist=True)
+        self.log("test/sw_loss", result["sw_loss"], on_epoch=True, sync_dist=True)
+        self.log("test/masked_mae", result["masked_mae"], on_epoch=True, sync_dist=True)
+        self.log("test/masked_corr", result["masked_corr"], on_epoch=True, sync_dist=True)
+        self.log("test/mask_rate", result["mask_rate"], on_epoch=True, sync_dist=True)
+        return {"test_loss": total_loss}
+
     def configure_optimizers(self):  # type: ignore[override]
         optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
         scheduler_dict = configure_scheduler(optimizer, self.scheduler_config)