pytorch lightning removal

app.py

@@ -1,20 +1,19 @@
 import operator
 import torch
-from torchvision import transforms
 import gradio as gr
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image
 import gradio as gr
-import model
 from data_loader import CIFAR_CLASS_LABELS, TEST_TRANSFORM
 import matplotlib
+from model import ResNet18
 matplotlib.use('agg')
 from matplotlib import pyplot as plt
 
-resnet_18 = model.LitResnet()
-state_dict = torch.load("resnet.pth", map_location=torch.device('cpu'))
-resnet_18.load_state_dict(state_dict)
-resnet_18_model = resnet_18.model
+
+resnet_18_model = ResNet18()
+resnet_18_model.load_state_dict(torch.load('resnet18.pth'))
+resnet_18_model.eval()
 
 classes = ('plane', 'car', 'bird', 'cat', 'deer',
            'dog', 'frog', 'horse', 'ship', 'truck')

model.py

@@ -79,121 +79,3 @@ class ResNet(nn.Module):
 def ResNet18():
     return ResNet(BasicBlock, [2, 2, 2, 2])
 
-import torch.nn as nn
-from torch.optim.lr_scheduler import OneCycleLR
-from torch.utils.data import DataLoader
-import matplotlib.pyplot as plt
-
-from data_loader import CifarAlbumentationsDataset,\
-    CIFAR_CLASS_LABELS, TRAIN_TRANSFORM, TEST_TRANSFORM
-import model
-from torch_lr_finder import LRFinder
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from pytorch_lightning import LightningModule
-from torch.optim.lr_scheduler import OneCycleLR
-from torchmetrics.functional import accuracy
-
-class LitResnet(LightningModule):
-    def __init__(self, lr=0.03, batch_size=512):
-        super().__init__()
-
-        self.save_hyperparameters()
-        self.criterion = nn.CrossEntropyLoss()
-        self.model = ResNet18()
-
-    def forward(self, x):
-        return self.model(x)
-
-    def training_step(self, batch, batch_idx):
-        x, y = batch
-        output = self.forward(x)
-        loss = self.criterion(output, y)
-        self.log("train_loss", loss)
-        acc = accuracy(torch.argmax(output, dim=1), 
-                       y, 'multiclass', num_classes=10)
-        self.log(f"train_acc", acc, prog_bar=True)
-        return loss
-
-    def evaluate(self, batch, stage=None):
-        x, y = batch
-        output = self.forward(x)
-        loss = self.criterion(output, y)
-        preds = torch.argmax(output, dim=1)
-        acc = accuracy(preds, y, 'multiclass', num_classes=10)
-
-        if stage:
-            self.log(f"{stage}_loss", loss, prog_bar=True)
-            self.log(f"{stage}_acc", acc, prog_bar=True)
-
-    def validation_step(self, batch, batch_idx):
-        self.evaluate(batch, "val")
-
-    def test_step(self, batch, batch_idx):
-        self.evaluate(batch, "test")
-
-    # todo
-    # change the default for num_iter
-    def lr_finder(self, optimizer,  num_iter=200,):
-        lr_finder = LRFinder(self, optimizer, self.criterion,
-            device=self.device)
-        lr_finder.range_test(
-            self.train_dataloader(), end_lr=1,
-            num_iter=num_iter, step_mode='exp',
-            )
-        ax, suggested_lr = lr_finder.plot(suggest_lr=True)
-        # todo
-        # how to log maplotlib images
-        # self.logger.experiment.add_image('lr_finder', plt.gcf(), 0)
-        lr_finder.reset() 
-        return suggested_lr
-    def configure_optimizers(self):
-        optimizer = torch.optim.SGD(
-            self.parameters(),
-            lr=self.hparams.lr,
-            momentum=0.9,
-            weight_decay=5e-4,
-        )
-        suggested_lr = self.lr_finder(optimizer)
-        steps_per_epoch = len(self.train_dataloader())
-        scheduler_dict = {
-            "scheduler":  OneCycleLR(
-                optimizer, max_lr=suggested_lr,
-                steps_per_epoch=steps_per_epoch,
-                epochs=self.trainer.max_epochs, 
-                pct_start=5/self.trainer.max_epochs,
-                three_phase=False,
-                div_factor=100,
-                final_div_factor=100,
-                anneal_strategy='linear',
-            ),
-            "interval": "step",
-        }
-        return {"optimizer": optimizer, "lr_scheduler": scheduler_dict}
-    ####################
-    # DATA RELATED HOOKS
-    ####################
-
-    def prepare_data(self, data_path='../data'):
-        CifarAlbumentationsDataset(
-                data_path, train=True, download=True)
-        CifarAlbumentationsDataset(
-                data_path, train=False, download=True)
-
-    def setup(self, stage=None, data_dir='../data'):
-
-        if stage == "fit" or stage is None:
-            self.train_dataset = CifarAlbumentationsDataset(data_dir, train=True, transform=TRAIN_TRANSFORM)
-            self.test_dataset = CifarAlbumentationsDataset(data_dir, train=False, transform=TEST_TRANSFORM)
-
-
-    def train_dataloader(self):
-        return DataLoader(self.train_dataset, batch_size=self.hparams.batch_size,
-        shuffle=True, pin_memory=True) #num_workers=4, 
-
-    def val_dataloader(self):
-        return DataLoader(self.test_dataset, batch_size=self.hparams.batch_size,
-        shuffle=False, pin_memory=True)
-

requirements.txt

@@ -5,4 +5,3 @@ grad-cam
 pillow
 numpy
 albumentations
-pytorch-lightning

resnet.pth → resnet18.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9bfd0a13c1f17b0977282a76b56aace7b682c1e3ac6ea12716ccbba97aa7afd3
-size 44774390
+oid sha256:f5cc1136d85d13f5a4152098208cdddf533cf344ffbfb30a06dcb698dddf862f
+size 44772860