fix: update accuracy for classification only

detector/model.py

@@ -150,27 +150,29 @@ class FontDetector(ptl.LightningModule):
         self.font_accur_train = torchmetrics.Accuracy(
             task="multiclass", num_classes=config.FONT_COUNT
         )
-        self.direction_accur_train = torchmetrics.Accuracy(
-            task="multiclass", num_classes=2
-        )
         self.font_accur_val = torchmetrics.Accuracy(
             task="multiclass", num_classes=config.FONT_COUNT
         )
-        self.direction_accur_val = torchmetrics.Accuracy(
-            task="multiclass", num_classes=2
-        )
         self.font_accur_test = torchmetrics.Accuracy(
             task="multiclass", num_classes=config.FONT_COUNT
         )
-        self.direction_accur_test = torchmetrics.Accuracy(
-            task="multiclass", num_classes=2
-        )
+        if not font_classification_only:
+            self.direction_accur_train = torchmetrics.Accuracy(
+                task="multiclass", num_classes=2
+            )
+            self.direction_accur_val = torchmetrics.Accuracy(
+                task="multiclass", num_classes=2
+            )
+            self.direction_accur_test = torchmetrics.Accuracy(
+                task="multiclass", num_classes=2
+            )
         self.lr = lr
         self.betas = betas
         self.num_warmup_iters = num_warmup_iters
         self.num_iters = num_iters
         self.num_epochs = num_epochs
         self.load_epoch = 0
+        self.font_classification_only = font_classification_only
 
     def forward(self, x):
         return self.model(x)
@@ -188,24 +190,26 @@ class FontDetector(ptl.LightningModule):
             self.font_accur_train(y_hat[..., : config.FONT_COUNT], y[..., 0]),
             sync_dist=True,
         )
-        self.log(
-            "train_direction_accur",
-            self.direction_accur_train(
-                y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1]
-            ),
-            sync_dist=True,
-        )
+        if not self.font_classification_only:
+            self.log(
+                "train_direction_accur",
+                self.direction_accur_train(
+                    y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1]
+                ),
+                sync_dist=True,
+            )
         return {"loss": loss}
 
     def on_train_epoch_end(self) -> None:
         self.log("train_font_accur", self.font_accur_train.compute(), sync_dist=True)
-        self.log(
-            "train_direction_accur",
-            self.direction_accur_train.compute(),
-            sync_dist=True,
-        )
         self.font_accur_train.reset()
-        self.direction_accur_train.reset()
+        if not self.font_classification_only:
+            self.log(
+                "train_direction_accur",
+                self.direction_accur_train.compute(),
+                sync_dist=True,
+            )
+            self.direction_accur_train.reset()
 
     def validation_step(
         self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
@@ -215,18 +219,22 @@ class FontDetector(ptl.LightningModule):
         loss = self.loss(y_hat, y)
         self.log("val_loss", loss, prog_bar=True, sync_dist=True)
         self.font_accur_val.update(y_hat[..., : config.FONT_COUNT], y[..., 0])
-        self.direction_accur_val.update(
-            y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1]
-        )
+        if not self.font_classification_only:
+            self.direction_accur_val.update(
+                y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1]
+            )
         return {"loss": loss}
 
     def on_validation_epoch_end(self):
         self.log("val_font_accur", self.font_accur_val.compute(), sync_dist=True)
-        self.log(
-            "val_direction_accur", self.direction_accur_val.compute(), sync_dist=True
-        )
         self.font_accur_val.reset()
-        self.direction_accur_val.reset()
+        if not self.font_classification_only:
+            self.log(
+                "val_direction_accur",
+                self.direction_accur_val.compute(),
+                sync_dist=True,
+            )
+            self.direction_accur_val.reset()
 
     def test_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int):
         X, y = batch
@@ -234,18 +242,22 @@ class FontDetector(ptl.LightningModule):
         loss = self.loss(y_hat, y)
         self.log("test_loss", loss, prog_bar=True, sync_dist=True)
         self.font_accur_test.update(y_hat[..., : config.FONT_COUNT], y[..., 0])
-        self.direction_accur_test.update(
-            y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1]
-        )
+        if not self.font_classification_only:
+            self.direction_accur_test.update(
+                y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1]
+            )
         return {"loss": loss}
 
     def on_test_epoch_end(self) -> None:
         self.log("test_font_accur", self.font_accur_test.compute(), sync_dist=True)
-        self.log(
-            "test_direction_accur", self.direction_accur_test.compute(), sync_dist=True
-        )
         self.font_accur_test.reset()
-        self.direction_accur_test.reset()
+        if not self.font_classification_only:
+            self.log(
+                "test_direction_accur",
+                self.direction_accur_test.compute(),
+                sync_dist=True,
+            )
+            self.direction_accur_test.reset()
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(