Fix model params & add accuracy metric

mask_detector/dataset.py

@@ -9,7 +9,7 @@
 
 
 class MaskDataset(Dataset):
-    def __init__(self, csv_file, image_size=100):
+    def __init__(self, csv_file, image_size=32):
         self.dataFrame = pd.read_csv(csv_file)
 
         self.transform = Compose([

mask_detector/mask_classifier.py

@@ -5,6 +5,7 @@
 from torch.utils.data import DataLoader, random_split
 from torch.optim import Adam
 from pytorch_lightning import LightningModule, Trainer, seed_everything
+from pytorch_lightning.metrics import Accuracy
 from pytorch_lightning.callbacks import ModelCheckpoint
 from sklearn.metrics import accuracy_score
 
@@ -16,43 +17,37 @@
 class MaskClassifier(LightningModule):
     def __init__(self, net, learning_rate=1e-3):
         super().__init__()
-        self.save_hyperparameters()
         self.net = net
+        self.learning_rate = learning_rate
+        self.accuracy = Accuracy()
 
     def forward(self, x):
         return self.net(x)
 
     def training_step(self, batch, batch_idx):
         x, y = batch
         out = self.net(x)
-
         loss = binary_cross_entropy(out, y)
 
-        self.log('train_loss', loss, on_epoch=True)
-
-        return loss
+        return {'loss': loss, 'accuracy': self.accuracy(out, y)}
 
     def validation_step(self, batch, batch_idx):
         x, y = batch
         out = self.net(x)
         loss = binary_cross_entropy(out, y)
 
-        self.log('valid_loss', loss, on_step=True)
+        return {'loss': loss, 'accuracy': self.accuracy(out, y)}
 
     def test_step(self, batch, batch_idx):
         x, y = batch
         out = self.net(x)
         loss = binary_cross_entropy(out, y)
 
-        _, out = torch.max(out, dim=1)
-        val_acc = accuracy_score(out.cpu(), y.cpu())
-        val_acc = torch.tensor(val_acc)
-
-        return {'test_loss': loss, 'test_acc': val_acc}
+        return {'loss': loss, 'accuracy': self.accuracy(out, y)}
 
     def configure_optimizers(self):
         # self.hparams available because we called self.save_hyperparameters()
-        return Adam(self.parameters(), lr=self.hparams.learning_rate)
+        return Adam(self.parameters(), lr=self.learning_rate)
 
 
 def cli_main():
@@ -65,9 +60,9 @@ def cli_main():
     ds_train, ds_validate, ds_test = train_val_test_split(
         dataset, train_ratio=0.8, validate_ratio=0.1, test_ratio=0.1)
 
-    train_loader = DataLoader(ds_train, batch_size=128)
-    val_loader = DataLoader(ds_validate, batch_size=128)
-    test_loader = DataLoader(ds_test, batch_size=128)
+    train_loader = DataLoader(ds_train, batch_size=32)
+    val_loader = DataLoader(ds_validate, batch_size=32)
+    test_loader = DataLoader(ds_test, batch_size=32)
 
     # ------------
     # model
@@ -78,12 +73,9 @@ def cli_main():
     # ------------
     # training
     # ------------
-    checkpoint_callback = ModelCheckpoint(
-        verbose=True,
-        monitor='test_acc',
-        mode='max'
-    )
-    trainer = Trainer(max_epochs=1, checkpoint_callback=checkpoint_callback)
+    gpus = 1 if torch.cuda.is_available() else 0
+
+    trainer = Trainer(gpus=gpus, max_epochs=3)
     trainer.fit(model, train_loader, val_loader)
 
     # ------------

mask_detector/model.py

@@ -17,16 +17,8 @@ def __init__(self):
             MaxPool2d(kernel_size=(2, 2))
         )
 
-        self.convLayers3 = Sequential(
-            Conv2d(64, 128, kernel_size=(3, 3), padding=(1, 1)),
-            ReLU(),
-            MaxPool2d(kernel_size=(2, 2))
-        )
-
         self.linearLayers = Sequential(
-            Linear(in_features=128*25*25, out_features=2048),
-            ReLU(),
-            Linear(in_features=2048, out_features=1024),
+            Linear(in_features=64*8*8, out_features=1024),
             ReLU(),
             Linear(in_features=1024, out_features=1),
             Softmax()
@@ -35,8 +27,7 @@ def __init__(self):
     def forward(self, x):
         x = self.convLayers1(x)
         x = self.convLayers2(x)
-        x = self.convLayers3(x)
-        x = x.view(-1, 128*25*25)
+        x = x.view(-1, 64*8*8)
         x = self.linearLayers(x)
 
         return x