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

BaseCNNArhitecture/AlexNetCustom.py

@@ -0,0 +1,329 @@
+# %% Импорт библиотек
+
+import time
+import torch
+import kagglehub
+import torch.nn as nn
+import pytorch_lightning as pl
+
+from torch import optim
+from torchmetrics import MetricCollection
+from torchvision import datasets, transforms
+from torchmetrics.classification import Accuracy
+from torch.utils.data import DataLoader, random_split
+
+# автоматическое отслеживание lr
+from pytorch_lightning.callbacks import LearningRateMonitor
+# ранние остановки
+from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+# сохранение весов
+from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
+# %% Параметры
+
+dataset_root = "./Datasets/ImageNet"
+batch_size = 64
+num_workers = 8
+num_epoch = 2
+
+# %% Трансформации
+
+train_transform = transforms.Compose([
+    transforms.RandomResizedCrop(224),
+    transforms.RandomHorizontalFlip(),
+
+    transforms.ToTensor(),
+
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+eval_transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+
+    transforms.ToTensor(),
+
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+# %% Скачивание ImageNet
+
+path = kagglehub.competition_download('imagenet-object-localization-challenge',
+                                      output_dir=dataset_root,
+                                      force_download=False)
+print("Path to competition files:", path)
+# %% Датасет
+
+full_train_dataset = datasets.ImageNet(
+    root=dataset_root,
+    split="train",
+    transform=train_transform
+)
+
+train_size = int(0.9 * len(full_train_dataset))
+val_size = len(full_train_dataset) - train_size
+
+train_dataset, val_dataset = random_split(
+    full_train_dataset,
+    [train_size, val_size],
+    generator=torch.Generator().manual_seed(42)
+)
+
+
+val_base_dataset = datasets.ImageNet(
+    root=dataset_root,
+    split="train",
+    transform=eval_transform
+)
+
+val_dataset.dataset = val_base_dataset
+
+test_dataset = datasets.ImageNet(
+    root=dataset_root,
+    split="val",
+    transform=eval_transform
+)
+
+train_loader = DataLoader(
+    dataset=train_dataset,
+    batch_size=batch_size,
+    shuffle=True,
+    num_workers=num_workers,
+    pin_memory=True,
+    persistent_workers=True
+)
+
+val_loader = DataLoader(
+    dataset=val_dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=num_workers,
+    pin_memory=True,
+    persistent_workers=True
+)
+
+test_loader = DataLoader(
+    dataset=test_dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=num_workers,
+    pin_memory=True,
+    persistent_workers=True
+)
+
+# %% AlexNet
+
+class AlexNetCustom(nn.Module):
+    def __init__(self,
+                 in_channels: int = 3,
+                 out_classes: int=1_000):
+        super().__init__()
+        
+        self._out_classes = out_classes
+        
+        self.feature = nn.Sequential(
+            nn.Conv2d(in_channels, 96, kernel_size=11, stride=4, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            
+            nn.Conv2d(96, 256, kernel_size=5, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            
+            nn.Conv2d(256, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            
+            nn.Conv2d(384, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            
+            nn.Conv2d(384, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2)
+        )
+        
+        # Добавлено в современной реализации
+        self.avgpool = nn.AdaptiveAvgPool2d((6, 6)) 
+        
+        self.mlp = nn.Sequential(
+            nn.Dropout(p=0.5),
+            nn.Linear(256*6*6, 4096),
+            nn.ReLU(inplace=True),
+            
+            nn.Dropout(p=0.5),
+            nn.Linear(4096, 4096),
+            nn.ReLU(inplace=True),
+            
+            nn.Linear(4096, self._out_classes)
+        )
+        
+    def forward(self, image):
+        x = self.feature(image)
+        # x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.mlp(x)
+        
+        return x
+    
+# %% Классификатор обучатор
+
+class ClassificationLightning(pl.LightningModule):
+    def __init__(self,
+                 model: nn.Module,
+                 num_classes: int,
+                 optimizer_class,
+                 optimizer_kwargs: dict,
+                 scheduler_class=None,
+                 scheduler_kwargs: dict | None = None,
+                 scheduler_monitor: str = "val/loss"):
+        super().__init__()
+        
+        
+        scheduler_kwargs = scheduler_kwargs or {}
+        
+        self.model = model
+        self.optimizer_class = optimizer_class
+        self.scheduler_class = scheduler_class
+        
+        num_parameters = sum(
+            p.numel()
+            for p in model.parameters()
+        )
+        
+        num_trainable_parameters = sum(
+            p.numel()
+            for p in model.parameters()
+            if p.requires_grad
+        )
+        
+        self.save_hyperparameters({
+            "model": model.__class__.__name__,
+            "num_classes": num_classes,
+        
+            "num_parameters": num_parameters,
+            "num_trainable_parameters": num_trainable_parameters,
+        
+            "optimizer": optimizer_class.__name__,
+            "optimizer_kwargs": optimizer_kwargs,
+        
+            "scheduler": (
+                scheduler_class.__name__
+                if scheduler_class is not None
+                else None
+            ),
+            "scheduler_kwargs": scheduler_kwargs,
+            "scheduler_monitor": scheduler_monitor
+        })
+        
+        
+        self.criterion = nn.CrossEntropyLoss()
+        self.train_metrics = MetricCollection({
+            "acc_top1": Accuracy(task="multiclass", 
+                                 num_classes=self.hparams.num_classes,
+                                 top_k=1),
+            "acc_top5": Accuracy(task="multiclass", 
+                                 num_classes=self.hparams.num_classes,
+                                 top_k=5),  
+        })
+        
+        self.val_metrics = self.train_metrics.clone()
+        self.test_metrics = self.train_metrics.clone()
+        
+    def forward(self, image):
+        return self.model(image)
+    
+    def _shared_step(self,
+                    sample,
+                    stage,
+                    metrics):
+        x, y = sample
+        x = x.to(torch.float32)
+        y = y.to(torch.int64)
+        
+        
+        logits = self(x)
+        loss = self.criterion(logits, y)
+        
+        metrics.update(logits, y)
+        
+        self.log(f"{stage}/loss", 
+                 loss, 
+                 prog_bar=True, 
+                 on_epoch=True, 
+                 on_step=False
+         )
+        
+        return loss
+    
+    def _log_metrics(self, stage, metrics):
+        computed = metrics.compute()
+        
+        for name, value in computed.items():
+            self.log(f"{stage}/{name}", value, prog_bar=True)
+        
+        metrics.reset()
+    
+    def on_train_epoch_start(self):
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        
+        self._epoch_start_time = time.perf_counter()
+    
+    def training_step(self, batch, batch_idx):
+        return self._shared_step(batch, "train", self.train_metrics)
+    
+    def on_train_epoch_end(self):
+        self._log_metrics("train", self.train_metrics)
+        
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        
+        epoch_time = time.perf_counter() - self._epoch_start_time
+        lr = self.optimizers().param_groups[0]["lr"]
+        
+        self.log("train/lr", lr, prog_bar=True)
+        self.log("train/epoch_time", epoch_time)
+    
+    def validation_step(self, batch, batch_idx):
+        return self._shared_step(batch, "val", self.val_metrics)
+    
+    def on_validation_epoch_end(self):
+        self._log_metrics("val", self.val_metrics)
+    
+    def test_step(self, batch, batch_idx):
+        return self._shared_step(batch, "test", self.test_metrics)
+    
+    def on_test_epoch_end(self):
+       self._log_metrics("test", self.test_metrics)
+    
+    def configure_optimizers(self):
+        optimizer = self.optimizer_class(
+            self.parameters(),
+            **self.hparams.optimizer_kwargs
+        )
+        
+        if self.scheduler_class is None:
+            return optimizer
+        
+        scheduler = self.scheduler_class(
+            optimizer,
+            **self.hparams.scheduler_kwargs
+        )
+        
+        scheduler_config = {
+            "scheduler": scheduler,
+            "interval": "epoch",
+            "frequency": 1
+        }
+        
+        if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
+            scheduler_config["monitor"] = self.hparams.scheduler_monitor
+        
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": scheduler_config
+        }

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')}")
+

MyNNExamplePL.py

@@ -178,6 +178,8 @@ test_loader = DataLoader(dataset=test_dataset,
 # %% Запуск обучения
 
 model = MyNNExample(model=MyTorchCNNModel())
+model.compile()
+
 callbacks = [
     ModelCheckpoint(dirpath="models",
                     filename="{epoch}_{val_accuracy:.4f}_{val_loss:.4f}",
@@ -206,10 +208,10 @@ trainer = pl.Trainer(
 trainer.fit(model=model, 
             train_dataloaders=train_loader,
             val_dataloaders=test_loader)
-        
+
 # %% Инференс модели
         
-checkpoint_path = "./models/6_0.9460_0.1643.ckpt"
+checkpoint_path = "./models/7_0.9720_0.0958.ckpt"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = MyNNExample.load_from_checkpoint(checkpoint_path,
                                          model=MyTorchCNNModel(output_size=10))
@@ -226,4 +228,29 @@ with torch.no_grad():
     
 print(f"Предсказание: {pred_label}")
 print(f"Истинная метка: {true_label}")
-print(f"Уверенность: {probs[0, pred_label].item():.2%}")
+print(f"Уверенность: {probs[0, pred_label].item():.2%}")
+
+# %% Matlab export
+
+pytorch_model = model.model
+pytorch_model.eval()
+pytorch_model = pytorch_model.cpu()
+
+dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)
+
+with torch.no_grad():
+    _ = pytorch_model(dummy_input)
+    
+traced_model = torch.jit.trace(pytorch_model, dummy_input)
+
+with torch.no_grad():
+    y_original = pytorch_model(dummy_input)
+    y_traced = traced_model(dummy_input)
+    
+print("Original:", y_original.shape)
+print("Traced:", y_traced.shape)
+print("Max difference:", torch.max(torch.abs(y_original - y_traced)).item())
+
+traced_model.save("./models/my_model_matlab.pt")
+
+