Upload model.py

model.py

@@ -0,0 +1,49 @@
+import lightning as L
+from lightning.pytorch.utilities.model_summary import ModelSummary
+
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+import torchmetrics
+from torchvision import transforms
+
+from typing import Optional
+
+class ClassifierModel(L.LightningModule):
+    
+    def __init__(self, model: nn.Module, image_size: int = 500, learning_rate: float = 1e-3, num_classes: int = 3,
+                 train_transform: Optional[transforms.Compose] = None, val_transform: Optional[transforms.Compose] = None) -> None:
+        super().__init__()
+        self.model = model
+        self.learning_rate = learning_rate
+        self.example_input_array = torch.Tensor(5, 3, image_size, image_size)
+        self.f1_score = torchmetrics.F1Score(task='multiclass', num_classes=num_classes)
+        self.train_transform = train_transform
+        self.val_transform = val_transform
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+    
+    def print_summary(self) -> None:
+        print(ModelSummary(self, max_depth=-1))
+
+    def configure_optimizers(self) -> torch.optim.Optimizer:
+        return torch.optim.Adam(params=self.model.parameters(), lr=self.learning_rate)
+    
+    def training_step(self, batch: tuple, batch_idx: int) -> float:
+        X, y = batch
+        y_pred = self(X)
+        loss = F.cross_entropy(y_pred, y)
+        self.log_dict({'Train loss': loss, f'Train F1 score': self.f1_score(y_pred, y)}, 
+                      on_step=False, on_epoch=True)
+        return loss
+    
+    def validation_step(self, batch: tuple, batch_idx: int) -> float:
+        X, y = batch
+        y_pred = self(X)
+        loss = F.cross_entropy(y_pred, y)
+        self.log_dict({'Validation loss': loss, f'Validation F1 score': self.f1_score(y_pred, y)}, 
+                      on_step=False, on_epoch=True)
+        return loss
+