Update app.py

app.py

@@ -22,97 +22,8 @@ from torchvision.datasets import CIFAR10
 PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
 BATCH_SIZE = 64
 
-class LitCustomResNet(LightningModule):
-    def __init__(self, data_dir=PATH_DATASETS, learning_rate=0.05):
-
-        super().__init__()
-
-        self.model = CustomResNet()
-        # Set our init args as class attributes
-        self.data_dir = data_dir
-        self.learning_rate = learning_rate
-
-        # Hardcode some dataset specific attributes
-        self.num_classes = 10
-        self.train_transform = transforms.Compose([
-                          transforms.RandomHorizontalFlip(),
-                          transforms.RandomCrop(32, padding=4),
-                          transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
-                          transforms.ToTensor(),
-                          transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
-                         ])
-        self.test_transform = transforms.Compose([
-                          transforms.ToTensor(),
-                          transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
-                         ])
-        
-        self.accuracy = Accuracy(task='multiclass', num_classes=10)
-
-
-    def forward(self, x):
-        x = self.model(x)
-        x = x.view(-1, 10)
-        return F.log_softmax(x, dim=1)
-
-    def training_step(self, batch, batch_idx):
-        x, y = batch
-        logits = self(x)
-        loss = F.nll_loss(logits, y)
-        self.log("train_loss", loss)
-        return loss
-
-    def validation_step(self, batch, batch_idx):
-        x, y = batch
-        logits = self(x)
-        loss = F.nll_loss(logits, y)
-        preds = torch.argmax(logits, dim=1)
-        self.accuracy(preds, y)
-
-        # Calling self.log will surface up scalars for you in TensorBoard
-        self.log("val_loss", loss, prog_bar=True)
-        self.log("val_acc", self.accuracy, prog_bar=True)
-        return loss
-
-    def test_step(self, batch, batch_idx):
-        # Here we just reuse the validation_step for testing
-        return self.validation_step(batch, batch_idx)
-
-    def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
-        return optimizer
-
-    ####################
-    # DATA RELATED HOOKS
-    ####################
-
-    def prepare_data(self):
-        # download
-        CIFAR10(root=self.data_dir, train=True, download=True)
-        CIFAR10(root=self.data_dir, train=False, download=True)
-
-    def setup(self, stage=None):
-
-        # Assign train/val datasets for use in dataloaders
-        if stage == "fit" or stage is None:
-            cifar10_full = CIFAR10(self.data_dir, train=True, transform=self.train_transform)
-            self.cifar10_train, self.cifar10_val = random_split(cifar10_full, [45000, 5000])
-
-        # Assign test dataset for use in dataloader(s)
-        if stage == "test" or stage is None:
-            self.cifar10_test = CIFAR10(self.data_dir, train=False, transform=self.test_transform)
-
-    def train_dataloader(self):
-        return DataLoader(self.cifar10_train, batch_size=BATCH_SIZE, num_workers=os.cpu_count())
-
-    def val_dataloader(self):
-        return DataLoader(self.cifar10_val, batch_size=BATCH_SIZE, num_workers=os.cpu_count())
-
-    def test_dataloader(self):
-        return DataLoader(self.cifar10_test, batch_size=BATCH_SIZE, num_workers=os.cpu_count())
-
-model = LitCustomResNet.load_from_checkpoint('custom_resnet.ckpt', strict=False)
-#lit_model = torch.load(checkpoint)
-#model = lit_model.model
+model = CustomResNet(input_size=32,learning_rate=0.001,num_classes=10,)
+model.load_state_dict(torch.load("custom_resnet_model.pth", map_location=torch.device('cpu')), strict=False)
 
 inv_normalize = transforms.Normalize(
     mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],