Добавлена реализация LeNet и Alexnet

BaseCNNArhitecture/LeNetCustom.py

@@ -0,0 +1,244 @@
+# %% Импорт библиотек
+
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+
+from PIL import Image
+from torch.utils.data import DataLoader
+
+from torchvision import datasets, transforms
+from torchmetrics.classification import Accuracy
+# %% LeeNet
+
+class LeNetCustom(nn.Module):
+    
+    def __init__(self, 
+                 in_channels: int=1,
+                 out_classes: int=10):
+        super().__init__()
+        
+        self._out_classes=out_classes
+        
+        self.feature = nn.Sequential(
+            nn.Conv2d(in_channels, 6, kernel_size=5),
+            # nn.Tanh(),
+            nn.ReLU(),
+            nn.AvgPool2d(kernel_size=2, stride=2),
+            
+            nn.Conv2d(6, 16, kernel_size=5),
+            # nn.Tanh(),
+            nn.ReLU(),
+            nn.AvgPool2d(kernel_size=2, stride=2)          
+        )
+        
+        
+        self.mlp = nn.Sequential(
+            nn.Linear(256, 120),
+            # nn.LazyLinear(120),
+            # nn.Tanh(),
+            nn.ReLU(),
+            nn.Linear(120, 84),
+            # nn.Tanh(),
+            nn.ReLU(),
+            nn.Linear(84, out_classes)
+        )
+        
+    def forward(self, image):
+        x = self.feature(image)
+        x = self.mlp(x.flatten(1))
+        
+        return x
+    
+class LeNetCustomLightning(pl.LightningModule):
+    def __init__(self,
+                 model: nn.Module,
+                 lr_rate: float=0.1):
+        super().__init__()
+        
+        self.save_hyperparameters(ignore=["model"])
+        self.model = model
+        
+        self.criterion = nn.CrossEntropyLoss()
+        self.train_metric = Accuracy(task="multiclass",
+                                     num_classes=model._out_classes)
+        self.val_metric = self.train_metric.clone()
+        self.test_metric = self.train_metric.clone()
+
+    def forward(self, image):
+        return self.model(image)
+    
+    def _shared_step(self, sample, stage, metric):
+        x, y = sample
+        
+        logits = self.forward(x.to(torch.float32))
+        preds = torch.argmax(logits, 1)
+        loss = self.criterion(logits, y.to(torch.int64))
+        
+        metric.update(preds, y)
+        
+        self.log(f"{stage}_loss", loss, prog_bar=True, 
+                 on_epoch=True, on_step=False)
+        
+        return loss
+    
+    def _log_metric(self, stage, metric):
+        computed = metric.compute()
+        
+        self.log(f"{stage}_acc", computed, prog_bar=True)
+        
+        metric.reset()
+    
+    def training_step(self, batch, batch_idx):
+        return self._shared_step(batch, "train", self.train_metric)
+    
+    def on_train_epoch_end(self):
+        self._log_metric("train", self.train_metric)
+    
+    def validation_step(self, batch, batch_idx):
+        return self._shared_step(batch, "val", self.val_metric)
+    
+    def on_validation_epoch_end(self):
+        self._log_metric("val", self.val_metric)
+    
+    def test_step(self, batch, batch_idx):
+        return self._shared_step(batch, "test", self.test_metric)
+    
+    def on_test_epoch_end(self):
+        self._log_metric("test", self.test_metric)
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.SGD(self.parameters(), 
+                                    lr=self.hparams.lr_rate)
+        
+        return optimizer
+
+# %% Загрузка набора данных
+
+batch_size = 64
+
+image_transforms = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.1307,), (0.3081,))
+])
+
+train_dataset = datasets.MNIST(
+    root="./Datasets",
+    train=True,
+    download=False, # Если скачен, то убрать
+    transform=image_transforms
+)
+
+test_dataset = datasets.MNIST(
+    root="./Datasets",
+    train=False,
+    download=False, # Если скачен, то убрать
+    transform=image_transforms    
+)
+
+train_loader = DataLoader(dataset=train_dataset,
+                          batch_size=batch_size,
+                          shuffle=True)
+
+test_loader = DataLoader(dataset=test_dataset,
+                         batch_size=batch_size,
+                         shuffle=False)
+
+# %% Запуск обучения
+
+num_epoch=10
+
+lenet = LeNetCustom(in_channels=1,
+                    out_classes=10)
+lenet_lightning = LeNetCustomLightning(model=lenet,
+                                       lr_rate=0.1)
+lenet_lightning.compile()
+
+trainer = pl.Trainer(max_epochs=num_epoch,
+                     accelerator="gpu",
+                     devices=1)
+
+trainer.fit(model=lenet_lightning,
+            train_dataloaders=train_loader,
+            val_dataloaders=test_loader)
+
+
+
+# %% Matlab Export
+
+model = lenet_lightning.model
+model.eval()
+model = model.cpu()
+
+dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)
+with torch.no_grad():
+    _ = model(dummy_input)
+    
+traced_model = torch.jit.trace(model, dummy_input)
+traced_model.save("../models/LeNetCustom.pt")
+
+# %% ONNX экспорт
+
+model = lenet_lightning.model.cpu().eval()
+
+dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)
+
+torch.onnx.export(
+    model,
+    dummy_input,
+    "../models/LeNetCustom.onnx",
+    input_names=["input"],
+    output_names=["logits"],
+    dynamic_axes={
+        "input": {0: "batch_size"},
+        "logits": {0: "batch_size"}
+    },
+    opset_version=17,
+    # Проверка/оптимизация графа
+    do_constant_folding=True
+)
+
+# %% Новый экспорт в ONNX
+
+model = lenet_lightning.model.cpu().eval()
+dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)
+
+onnx_program = torch.onnx.export(
+    model,
+    (dummy_input,),
+    dynamo=True,
+    verify=True
+)
+
+onnx_program.save("../models/LeNetCustomNewExport.onnx")
+
+# %% Инференс модели
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = torch.jit.load("./models/LeNetCustom.pt", map_location=device)
+model.eval()
+
+image_transforms = transforms.Compose([
+    transforms.Grayscale(),
+    transforms.Resize((28,28)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.1307,), (0.3081,))
+])
+
+def predict(image_path):
+    image = Image.open(image_path)
+    image = image_transforms(image)
+    image = image.unsqueeze(0)
+    image = image.to(device)
+    
+    with torch.no_grad():
+        output = model(image)
+        
+        pred = output.argmax(dim=1).item()
+        
+    return pred
+
+print(f"Predicted label: {predict('./test_images/mnist/3.png')}")
+print(f"Predicted label: {predict('./test_images/mnist/5.png')}")
+print(f"Predicted label: {predict('./test_images/mnist/6.png')}")
+