Create model/network.py

model/network.py

@@ -0,0 +1,138 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+import pytorch_lightning as pl
+import torchmetrics
+from torch.optim.lr_scheduler import OneCycleLR
+from torchmetrics.functional import accuracy
+
+
+class ResBlock(nn.Module):
+
+  def __init__(self, in_channel, out_channel, stride=1):
+    super(ResBlock, self).__init__()
+    self.conv = nn.Sequential(
+        nn.Conv2d(in_channel, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+        nn.BatchNorm2d(in_channel),
+        nn.ReLU(),
+
+        nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=1, padding=1, bias=False),
+        nn.BatchNorm2d(out_channel),
+        nn.ReLU(),
+    )
+
+  def forward(self, x):
+    return(self.conv(x))
+
+
+
+class ResNet18(pl.LightningModule):
+  def __init__(self, train_loader_len, criterion, num_classes=10, lr=0.001, max_lr=1.45E-03):
+    super().__init__()
+    self.save_hyperparameters(ignore=['criterion'])
+
+    self.criterion = criterion
+    self.train_loader_len = train_loader_len
+    self.accuracy = torchmetrics.Accuracy(task="multiclass", num_classes=self.hparams.num_classes)
+
+    self.prep_layer = nn.Sequential(
+        nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False),
+        nn.BatchNorm2d(64),
+        nn.ReLU()
+    )
+
+    self.layer_one = nn.Sequential(
+        nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=False),
+        nn.MaxPool2d(2,2),
+        nn.BatchNorm2d(128),
+        nn.ReLU()
+    )
+
+    self.res_block1 = ResBlock(128, 128)
+
+    self.layer_two = nn.Sequential(
+        nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=False),
+        nn.MaxPool2d(2,2),
+        nn.BatchNorm2d(256),
+        nn.ReLU()
+    )
+
+    self.layer_three = nn.Sequential(
+        nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=False),
+        nn.MaxPool2d(2,2),
+        nn.BatchNorm2d(512),
+        nn.ReLU()
+    )
+
+    self.res_block2 = ResBlock(512, 512)
+
+    self.max_pool = nn.MaxPool2d(4,4)
+    self.fc = nn.Linear(512, num_classes, bias=False)
+
+  def forward(self, x):
+    x = self.prep_layer(x)
+    x = self.layer_one(x)
+    R1 = self.res_block1(x)
+    x = x + R1
+
+    x = self.layer_two(x)
+
+    x = self.layer_three(x)
+    R2 = self.res_block2(x)
+    x = x + R2
+
+    x = self.max_pool(x)
+
+    x = x.view(x.size(0), -1)
+    x = self.fc(x)
+
+    return(x)
+
+  def configure_optimizers(self):
+    optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr, weight_decay=1e-4)
+    scheduler = OneCycleLR(
+                            optimizer, 
+                            max_lr=self.hparams.max_lr,
+                            epochs=self.trainer.max_epochs,
+                            steps_per_epoch=self.train_loader_len,
+                            pct_start=5/self.trainer.max_epochs,
+                            div_factor=100,
+                            three_phase=False,
+                        )
+    if self.hparams.max_lr==1.45E-03:
+      return(optimizer)
+    else:
+      return([optimizer], [scheduler])
+
+  def training_step(self, train_batch, batch_idx):
+    data, target = train_batch
+    y_pred = self(data)
+    loss = self.criterion(y_pred, target)
+
+    pred = torch.argmax(y_pred.squeeze(), dim=1)
+    acc = accuracy(pred, target, task="multiclass", num_classes=self.hparams.num_classes)
+
+    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 validation_step(self, batch, batch_idx):
+    return(self.evaluate(batch, 'val'))
+
+  def test_step(self, batch, batch_idx):
+    return(self.evaluate(batch, 'test'))
+
+  def evaluate(self, batch, stage=None):
+      data, target = batch
+      y_pred = self(data)
+
+      loss = self.criterion(y_pred, target).item()
+      pred = torch.argmax(y_pred.squeeze(), dim=1)
+      acc = accuracy(pred, target, task="multiclass", num_classes=self.hparams.num_classes)
+
+      if stage:
+          self.log(f"{stage}_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
+          self.log(f"{stage}_acc", acc, prog_bar=True, on_step=False, on_epoch=True)
+
+      return pred, target