diff --git a/BaseCNNArhitecture/AlexNetCustom.py b/BaseCNNArhitecture/AlexNetCustom.py new file mode 100644 index 0000000..2796800 --- /dev/null +++ b/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 + } diff --git a/BaseCNNArhitecture/LeNetCustom.py b/BaseCNNArhitecture/LeNetCustom.py new file mode 100644 index 0000000..b4d158b --- /dev/null +++ b/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')}") + diff --git a/MyNNExamplePL.py b/MyNNExamplePL.py index 3d2920f..087716b 100644 --- a/MyNNExamplePL.py +++ b/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%}") \ No newline at end of file +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") + +