update lightning_model

lightningmodel.py

@@ -2,6 +2,7 @@ import os
 import math
 import numpy as np
 import pandas as pd
+import seaborn as sn
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -24,6 +25,7 @@ from PIL import Image
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image
 
+
 seed_everything(7)
 
 
@@ -109,9 +111,9 @@ class Net_S13(nn.Module):
         O = 10
         self.lastLayer = nn.Linear(I, O)
 
-        self.aGAP = nn.AdaptiveAvgPool2d((1, 1))
-        self.flat = nn.Flatten(1, -1)
-        self.gap = nn.AvgPool2d(avg)
+        # self.aGAP = nn.AdaptiveAvgPool2d((1, 1))
+        # self.flat = nn.Flatten(1, -1)
+        # self.gap = nn.AvgPool2d(avg)
         self.drop = nn.Dropout(drop)
 
     # convolution Block
@@ -186,11 +188,10 @@ class Net_S13(nn.Module):
         self.printf(4.0, x, "pool input")
         x = self.pool(x)
         self.printf(4.1, x, "pool output")
-        self.printEmpty()
-
-        # x = x.view(-1, 10)
-        self.printf(4.2, x, "For showing before last layer")
         x = x.view(x.size(0), -1)
+        self.printf(4.2, x, "after view")
+        self.printEmpty()
+        
         self.printf(5.0, x, "last layer input") #512, 1, 1
         x = self.lastLayer(x)
         # x = self.gap(x)
@@ -242,6 +243,7 @@ class LitResnet(LightningModule):
         self.evaluate(batch, "test")
 
     def configure_optimizers(self):
+        BATCH_SIZE = 256 if torch.cuda.is_available() else 64
         optimizer = torch.optim.SGD(
             self.parameters(),
             lr=self.hparams.lr,