Update model.py

model.py

@@ -31,6 +31,7 @@ class ResBlock(nn.Module):
         out = self.relu(out)
         return out
 
+
 class LightningDavidNet(LightningModule):
 
     def __init__(self,data_dir=PATH_DATASETS, hidden_size=16, learning_rate=2e-4,kernel_size=3, stride=1, padding=1, downsample = None):
@@ -75,7 +76,6 @@ class LightningDavidNet(LightningModule):
         x = self.r2(x)
         x=residual+x
         x = self.maxPool(x)
-        # # x = self.avgpool(x)
         x = x.view(-1,512)
         x = self.fc1(x)
         x = F.log_softmax(x, dim=1)
@@ -83,41 +83,40 @@ class LightningDavidNet(LightningModule):
 
     def training_step(self, batch, batch_idx):
         x,y = batch
-        loss = F.cross_entropy(self(x), y)
-        self.log("train_loss", loss)
+        y_pred = self(x)
+        loss = F.cross_entropy(y_pred, y)
+        acc  = self.accuracy(y_pred, y)
+
+        self.log('train_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
+        self.log('train_acc', acc, prog_bar=True, on_step=False, on_epoch=True)
+
         return loss
+        
+    def evaluate(self, batch, stage=None):
+        x, y = batch
+        y_test_pred = self(x)
+        loss = F.cross_entropy(y_test_pred, y)
+        acc  = self.accuracy(y_test_pred, y)
+
+        if stage:
+            self.log(f"{stage}_loss", loss, prog_bar=True)
+            self.log(f"{stage}_acc", acc, prog_bar=True)
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=0.03, weight_decay=1e-4)
-        steps_per_epoch = len(train_loader)
-        scheduler_dict = {
-            "scheduler": torch.optim.lr_scheduler.OneCycleLR(
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(
                 optimizer,
-                0.1,
+                max_lr= 5.38E-02, #self.hparams.lr,
+                pct_start = 5/self.trainer.max_epochs,
                 epochs=self.trainer.max_epochs,
-                steps_per_epoch=steps_per_epoch,
-            ),
-            "interval": "step",
-        }
-        return {"optimizer": optimizer, "lr_scheduler": scheduler_dict}
-
-        # lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, step_size=1)
-        # return [optimizer], [lr_scheduler]
-        # return optimizer
+                steps_per_epoch=len(train_loader),
+                div_factor=100,verbose=False,
+                three_phase=False
+            )
+        return ([optimizer],[scheduler])
 
     def validation_step(self, batch, batch_idx):
-        x,y = batch
-        logits = self(x)
-        loss = F.cross_entropy(logits, y)
-        preds = torch.argmax(logits,dim = 1)
-        self.accuracy(preds,y)
-        self.log("val_loss",loss, prog_bar = True)
-        self.log("val_arr",self.accuracy,prog_bar = True)
+        self.evaluate(batch, "val")
 
     def test_step(self,batch,batch_idx):
-        return self.validation_step(batch,batch_idx)
-
-    def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        x,y = batch
-        output = self(x)
-        return x,y,output.argmax(dim=1),output
+        self.evaluate(batch, "test")