From 056dd877469424edebe21c4320667566d1134bd2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=93=D1=80=D0=B0?= =?UTF-8?q?=D0=B1=D0=B0=D1=80=D1=8C?= Date: Mon, 28 Sep 2026 15:40:05 +1000 Subject: [PATCH] =?UTF-8?q?=D0=9E=D0=B1=D0=BD=D0=B0=D0=B2=D0=BB=D0=B5?= =?UTF-8?q?=D0=BD=D0=B0=20=D1=81=D1=82=D1=80=D1=83=D0=BA=D1=82=D1=83=D1=80?= =?UTF-8?q?=D0=B0=20=D0=BF=D1=80=D0=BE=D0=B5=D0=BA=D1=82=D0=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 2 + .gitignore | 8 +- BaseCNNArhitecture/AlexNetCustom.py | 488 ------------------ BaseCNNArhitecture/LeNetCustom.py | 244 --------- custom_arhitecture/__init__.py | 9 + custom_arhitecture/alexnet.py | 53 ++ custom_arhitecture/lenet.py | 38 ++ custom_arhitecture/vgg.py | 84 +++ datasets/__init__.py | 7 + datasets/imagenet.py | 158 ++++++ .../MyDataExamplePL.py | 0 MyNNExamplePL.py => exampels/MyNNExamplePL.py | 0 SimpleMLP.py => exampels/SimpleMLP.py | 0 ffn.py => exampels/ffn.py | 0 main.py | 142 +++++ train/__init__.py | 6 + train/classification.py | 165 ++++++ 17 files changed, 671 insertions(+), 733 deletions(-) create mode 100644 .env.example delete mode 100644 BaseCNNArhitecture/AlexNetCustom.py delete mode 100644 BaseCNNArhitecture/LeNetCustom.py create mode 100644 custom_arhitecture/__init__.py create mode 100644 custom_arhitecture/alexnet.py create mode 100644 custom_arhitecture/lenet.py create mode 100644 custom_arhitecture/vgg.py create mode 100644 datasets/__init__.py create mode 100644 datasets/imagenet.py rename MyDataExamplePL.py => exampels/MyDataExamplePL.py (100%) rename MyNNExamplePL.py => exampels/MyNNExamplePL.py (100%) rename SimpleMLP.py => exampels/SimpleMLP.py (100%) rename ffn.py => exampels/ffn.py (100%) create mode 100644 main.py create mode 100644 train/__init__.py create mode 100644 train/classification.py diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..8f9a273 --- /dev/null +++ b/.env.example @@ -0,0 +1,2 @@ +EXPERIMENT_NAME= +TRACKING_URI= \ No newline at end of file diff --git a/.gitignore b/.gitignore index 0bb651d..d18c5d3 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,10 @@ +.vscode/ +__pycache__/ + models/ Datasets/ test_images/ -train_logs/ \ No newline at end of file +train_logs/ +lightning_logs/ + +.env \ No newline at end of file diff --git a/BaseCNNArhitecture/AlexNetCustom.py b/BaseCNNArhitecture/AlexNetCustom.py deleted file mode 100644 index 4a1bff2..0000000 --- a/BaseCNNArhitecture/AlexNetCustom.py +++ /dev/null @@ -1,488 +0,0 @@ -# %% Импорт библиотек - -import time -import torch -import kagglehub -import pandas as pd -import torch.nn as nn -import pytorch_lightning as pl - -from PIL import Image -from torch import optim -from pathlib import Path -from torchmetrics import MetricCollection -from typing import Optional, Callable, Any -from torchvision import datasets, transforms -from torchmetrics.classification import Accuracy -from torch.utils.data import DataLoader, Dataset - -# Логирование данных -from pytorch_lightning.loggers import MLFlowLogger -# автоматическое отслеживание 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 = "/media/daniel/developments/Datasets/ImageNet" -batch_size = 128 -num_workers = 8 -num_epoch = 90 -learning_rate = 3e-4 - -# %% Трансформации - -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) -# %% Датасет - -TransformType = Optional[Callable[..., Any]] - -class ImageNetValDataset(Dataset): - def __init__( - self, - image_dir, - labels_file, - class_to_idx, - classes, - transform=None - ): - self.image_dir = Path(image_dir) - self.transform = transform - - self.class_to_idx = class_to_idx - self.classes = classes - - df = pd.read_csv(labels_file) - - self.samples = [] - - for _, row in df.iterrows(): - image_id = row["ImageId"] - synset = row["PredictionString"].split()[0] - label = self.class_to_idx[synset] - image_path = self.image_dir / f"{image_id}.JPEG" - self.samples.append( - (image_path, label) - ) - - def __len__(self): - return len(self.samples) - - def __getitem__(self, idx): - image_path, label = self.samples[idx] - - image = Image.open(image_path).convert("RGB") - - if self.transform is not None: - image = self.transform(image) - - return image, label - - -class ImageNetDataModule(pl.LightningDataModule): - def __init__(self, - dataset_root: Path, - train_transform: TransformType, - eval_transform: TransformType, - batch_size: int = 64, - num_workers:int = 8 - ): - super().__init__() - - self.dataset_root = Path(dataset_root) - - self.train_transform = train_transform - self.eval_transform = eval_transform - self.num_workers = num_workers - self.batch_size = batch_size - - self.imagenet_root = Path(dataset_root) / "ILSVRC" - - self.train_dir = self.imagenet_root / "Data" / "CLS-LOC" / "train" - self.val_dir = self.imagenet_root / "Data" / "CLS-LOC" / "val" - - self.mapping_file = Path(dataset_root) / "LOC_synset_mapping.txt" - self.val_labels_file = Path(dataset_root) / "LOC_val_solution.csv" - - self.train_dataset = None - self.val_dataset = None - - self.class_to_idx = None - self.classes = None - - def setup(self, stage=None): - synset_to_name = {} - with open(self.mapping_file, "r", encoding="utf-8") as f: - for line in f: - synset, class_name = line.strip().split(" ", 1) - synset_to_name[synset] = class_name - - train_dataset = datasets.ImageFolder(root=self.train_dir, - transform=self.train_transform) - - self.class_to_idx = train_dataset.class_to_idx.copy() - self.classes = [synset_to_name[synset] - for synset in train_dataset.classes] - train_dataset.classes = self.classes - self.train_dataset = train_dataset - - self.val_dataset = ImageNetValDataset( - image_dir=self.val_dir, - labels_file=self.val_labels_file, - class_to_idx=self.class_to_idx, - classes=self.classes, - transform=self.eval_transform - ) - - def train_dataloader(self): - return DataLoader( - dataset=self.train_dataset, - batch_size=self.batch_size, - shuffle=True, - num_workers=self.num_workers, - pin_memory=True, - persistent_workers=self.num_workers > 0, - prefetch_factor=4 if self.num_workers > 0 else None - ) - - def val_dataloader(self): - return DataLoader( - dataset=self.val_dataset, - batch_size=self.batch_size, - shuffle=False, - num_workers=self.num_workers, - pin_memory=True, - persistent_workers=self.num_workers > 0, - prefetch_factor=4 if self.num_workers > 0 else None - ) - - def test_dataloader(self): - return DataLoader( - dataset=self.val_dataset, - batch_size=self.batch_size, - shuffle=False, - num_workers=self.num_workers, - pin_memory=True, - persistent_workers=self.num_workers > 0 - ) - -# %% 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 - - 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 - } - -# %% Обучение - -pl.seed_everything(42, workers=True) -torch.set_float32_matmul_precision("high") - -imagenet_dm = ImageNetDataModule( - dataset_root=dataset_root, - train_transform=train_transform, - eval_transform=eval_transform, - batch_size=batch_size, - num_workers=num_workers -) - - -optimizer_Adam = optim.Adam -optimizer_Adam_kwargs = { - "lr": 3e-4, - "weight_decay": 1e-4 -} - -scheduler_ReduceLR = optim.lr_scheduler.ReduceLROnPlateau -scheduler_ReduceLR_kwargs={ - "mode": "min", - "factor": 0.1, - "patience": 5 -} - -callbacks = [ - ModelCheckpoint( - dirpath="models", - filename="epoch={epoch}", - auto_insert_metric_name=False, - save_last=True, - save_top_k=1, - monitor="val/acc_top1", - mode="max" - ), - - LearningRateMonitor( - logging_interval="epoch" - ), - - EarlyStopping( - monitor="val/loss", - min_delta=2e-4, - patience=10, - verbose=False, - mode="min" - ) -] - -logger = MLFlowLogger( - experiment_name="ImageNetMyAlgoritms", - run_name="AlexNetCustom", - tracking_uri="http://10.16.18.1:8055", - log_model=True, -) - -alexnet = AlexNetCustom(in_channels=3, out_classes=1000) -model_adam = ClassificationLightning( - model=alexnet, - num_classes=1000, - optimizer_class=optimizer_Adam, - optimizer_kwargs=optimizer_Adam_kwargs, - scheduler_class=scheduler_ReduceLR, - scheduler_kwargs=scheduler_ReduceLR_kwargs, - scheduler_monitor="val/loss" -) - -trainer = pl.Trainer( - max_epochs=num_epoch, - accelerator="gpu", - devices=1, - logger=logger, - # precision="16-mixed", - callbacks=callbacks -) - -trainer.fit(model=model_adam, - datamodule=imagenet_dm) diff --git a/BaseCNNArhitecture/LeNetCustom.py b/BaseCNNArhitecture/LeNetCustom.py deleted file mode 100644 index b4d158b..0000000 --- a/BaseCNNArhitecture/LeNetCustom.py +++ /dev/null @@ -1,244 +0,0 @@ -# %% Импорт библиотек - -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/custom_arhitecture/__init__.py b/custom_arhitecture/__init__.py new file mode 100644 index 0000000..7a2e47e --- /dev/null +++ b/custom_arhitecture/__init__.py @@ -0,0 +1,9 @@ +from .lenet import LeNetCustom +from .alexnet import AlexNetCustom +from .vgg import VGGCustom + +__all__ = [ + 'LeNetCustom', + 'AlexNetCustom', + 'VGGCustom' +] \ No newline at end of file diff --git a/custom_arhitecture/alexnet.py b/custom_arhitecture/alexnet.py new file mode 100644 index 0000000..6070560 --- /dev/null +++ b/custom_arhitecture/alexnet.py @@ -0,0 +1,53 @@ +import torch +import torch.nn as nn + + +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 diff --git a/custom_arhitecture/lenet.py b/custom_arhitecture/lenet.py new file mode 100644 index 0000000..a20f658 --- /dev/null +++ b/custom_arhitecture/lenet.py @@ -0,0 +1,38 @@ +import torch.nn as nn + + +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.AvgPool2d(kernel_size=2, stride=2), + + nn.Conv2d(6, 16, kernel_size=5), + nn.Tanh(), + nn.AvgPool2d(kernel_size=2, stride=2), + + nn.Conv2d(16, 120, kernel_size=5), + nn.Tanh(), + ) + + + self.mlp = nn.Sequential( + nn.LazyLinear(84), + nn.Tanh(), + nn.Linear(84, out_classes) + ) + + def forward(self, image): + x = self.feature(image) + x = self.mlp(x.flatten(1)) + + return x + diff --git a/custom_arhitecture/vgg.py b/custom_arhitecture/vgg.py new file mode 100644 index 0000000..86a346b --- /dev/null +++ b/custom_arhitecture/vgg.py @@ -0,0 +1,84 @@ +import torch +import torch.nn as nn + + +VGG_CONFIG = { + "VGG16": [64, 64, "M", + 128, 128, "M", + 256, 256, 256, "M", + 512, 512, 512, "M", + 512, 512, 512, "M"], + "VGG19": [64, 64, "M", + 128, 128, "M", + 256, 256, 256, 256, "M", + 512, 512, 512, 512, "M", + 512, 512, 512, 512, "M"] +} + + +class VGGCustom(nn.Module): + def __init__(self, + architecture:str="VGG16", + num_channels:int = 3, + num_classes:int=10, + batch_normalization=False, + dropout:float=0.5) -> None: + + super().__init__() + + if architecture not in VGG_CONFIG: + raise ValueError(f"Architecture {architecture} is invalid. ", + f"Available architectures are :{list(VGG_CONFIG.keys())}") + + + self.batch_normalization = batch_normalization + + self.features = self._make_layers( + VGG_CONFIG[architecture], + in_channels=num_channels + ) + + self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) + + self.classifier = nn.Sequential( + nn.Linear(512*7*7, 4096), + nn.ReLU(True), + nn.Dropout(p=dropout), + + nn.Linear(4096, 4096), + nn.ReLU(True), + nn.Dropout(p=dropout), + + nn.Linear(4096, num_classes) + ) + + def _make_layers(self, cfg, in_channels): + layers = [] + + for layer in cfg: + if layer == 'M': + layers.append(nn.MaxPool2d(kernel_size=2, stride=2)) + else: + + layers.append(nn.Conv2d(in_channels=in_channels, + out_channels=layer, + kernel_size=3, + stride=1, + padding=1)) + + if self.batch_normalization: + layers.append(nn.BatchNorm2d(layer)) + + layers.append(nn.ReLU(True)) + + in_channels = layer + + return nn.Sequential(*layers) + + def forward(self, image): + x = self.features(image) + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.classifier(x) + + return x \ No newline at end of file diff --git a/datasets/__init__.py b/datasets/__init__.py new file mode 100644 index 0000000..f6dfe69 --- /dev/null +++ b/datasets/__init__.py @@ -0,0 +1,7 @@ +from .imagenet import ImageNetDataModule, ImageNetValDataset + + +__all__ = [ + "ImageNetDataModule", + "ImageNetValDataset", +] \ No newline at end of file diff --git a/datasets/imagenet.py b/datasets/imagenet.py new file mode 100644 index 0000000..22931fd --- /dev/null +++ b/datasets/imagenet.py @@ -0,0 +1,158 @@ +import torch +import pandas as pd +import pytorch_lightning as pl + +from PIL import Image +from pathlib import Path +from torchvision import datasets +from torchvision.io import read_file, decode_image, ImageReadMode +from typing import Optional, Callable, Any +from torch.utils.data import DataLoader, Dataset + + +TransformType = Optional[Callable[..., Any]] + + +def image_loader(path: str) -> torch.Tensor: + data = read_file(path) + + return decode_image( + data, + mode=ImageReadMode.RGB + ) + + +class ImageNetValDataset(Dataset): + def __init__( + self, + image_dir, + labels_file, + class_to_idx, + classes, + transform=None + ): + self.image_dir = Path(image_dir) + self.transform = transform + + self.class_to_idx = class_to_idx + self.classes = classes + + df = pd.read_csv( + labels_file, + usecols=["ImageId", "PredictionString"] + ) + + image_ids = df["ImageId"].to_numpy() + + synsets = ( + df["PredictionString"] + .str.split() + .str[0] + .to_numpy() + ) + + self.samples = [ + ( + self.image_dir / f"{image_id}.JPEG", + class_to_idx[synset] + ) + for image_id, synset in zip(image_ids, synsets) + ] + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx): + image_path, label = self.samples[idx] + + image = decode_image( + str(image_path), + mode=ImageReadMode.RGB + ) + + if self.transform is not None: + image = self.transform(image) + + return image, label + + +class ImageNetDataModule(pl.LightningDataModule): + def __init__(self, + dataset_root: Path, + train_transform: TransformType, + eval_transform: TransformType, + batch_size: int = 64, + num_workers:int = 8 + ): + super().__init__() + + self.dataset_root = Path(dataset_root) + + self.train_transform = train_transform + self.eval_transform = eval_transform + self.num_workers = num_workers + self.batch_size = batch_size + + self.imagenet_root = Path(dataset_root) / "ILSVRC" + + self.train_dir = self.imagenet_root / "Data" / "CLS-LOC" / "train" + self.val_dir = self.imagenet_root / "Data" / "CLS-LOC" / "val" + + self.mapping_file = Path(dataset_root) / "LOC_synset_mapping.txt" + self.val_labels_file = Path(dataset_root) / "LOC_val_solution.csv" + + self.train_dataset = None + self.val_dataset = None + + self.class_to_idx = None + self.classes = None + + def setup(self, stage=None): + synset_to_name = {} + with open(self.mapping_file, "r", encoding="utf-8") as f: + for line in f: + synset, class_name = line.strip().split(" ", 1) + synset_to_name[synset] = class_name + + train_dataset = datasets.ImageFolder(root=self.train_dir, + transform=self.train_transform, + loader=image_loader) + + self.class_to_idx = train_dataset.class_to_idx.copy() + self.classes = [synset_to_name[synset] + for synset in train_dataset.classes] + train_dataset.classes = self.classes + self.train_dataset = train_dataset + + self.val_dataset = ImageNetValDataset( + image_dir=self.val_dir, + labels_file=self.val_labels_file, + class_to_idx=self.class_to_idx, + classes=self.classes, + transform=self.eval_transform + ) + + def train_dataloader(self): + return DataLoader( + dataset=self.train_dataset, + batch_size=self.batch_size, + shuffle=True, + num_workers=self.num_workers, + pin_memory=True, + persistent_workers=self.num_workers > 0, + prefetch_factor=2 if self.num_workers > 0 else None, + drop_last=True + ) + + def val_dataloader(self): + return DataLoader( + dataset=self.val_dataset, + batch_size=self.batch_size, + shuffle=False, + num_workers=self.num_workers, + pin_memory=True, + persistent_workers=self.num_workers > 0, + prefetch_factor=2 if self.num_workers > 0 else None, + drop_last=False + ) + diff --git a/MyDataExamplePL.py b/exampels/MyDataExamplePL.py similarity index 100% rename from MyDataExamplePL.py rename to exampels/MyDataExamplePL.py diff --git a/MyNNExamplePL.py b/exampels/MyNNExamplePL.py similarity index 100% rename from MyNNExamplePL.py rename to exampels/MyNNExamplePL.py diff --git a/SimpleMLP.py b/exampels/SimpleMLP.py similarity index 100% rename from SimpleMLP.py rename to exampels/SimpleMLP.py diff --git a/ffn.py b/exampels/ffn.py similarity index 100% rename from ffn.py rename to exampels/ffn.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..7c8347f --- /dev/null +++ b/main.py @@ -0,0 +1,142 @@ +import os +import torch +import mlflow +import pytorch_lightning as pl + +from torch import optim +from dotenv import load_dotenv +from torchvision.transforms import v2 +from datasets import ImageNetDataModule +from train import ClassificationLightning +from custom_arhitecture import VGGCustom + +from pytorch_lightning.loggers import MLFlowLogger +from pytorch_lightning.callbacks import LearningRateMonitor +from pytorch_lightning.callbacks.early_stopping import EarlyStopping +from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint + + +load_dotenv() + +IMAGENET_MEAN = [0.485, 0.456, 0.406] +IMAGENET_STD = [0.229, 0.224, 0.225] + +dataset_root = "/media/daniel/developments/Datasets/ImageNet" +batch_size = 48 +num_workers = 8 +num_epoch = 90 +learning_rate = 3e-4 + +train_transform = v2.Compose([ + v2.RandomResizedCrop( + size=(224, 224), + scale=(0.08, 1.0), + ratio=(3 / 4, 4 / 3), + antialias=True + ), + v2.RandomHorizontalFlip(p=0.5), + v2.ToDtype(torch.float32, scale=True), + v2.Normalize( + mean=IMAGENET_MEAN, + std=IMAGENET_STD + ) +]) + +eval_transform = v2.Compose([ + v2.Resize(256, antialias=True), + v2.CenterCrop(224), + v2.ToDtype(torch.float32, scale=True), + v2.Normalize( + mean=IMAGENET_MEAN, + std=IMAGENET_STD + ) +]) + + +pl.seed_everything(42, workers=True) +torch.set_float32_matmul_precision("high") + +imagenet_dm = ImageNetDataModule( + dataset_root=dataset_root, + train_transform=train_transform, + eval_transform=eval_transform, + batch_size=batch_size, + num_workers=num_workers +) + +optimizer_Adam_kwargs = { + "lr": 3e-4, + "weight_decay": 1e-4 +} + +scheduler_ReduceLR_kwargs={ + "mode": "min", + "factor": 0.1, + "patience": 5 +} + +callbacks = [ + ModelCheckpoint( + dirpath="models", + filename="epoch={epoch}", + auto_insert_metric_name=False, + save_last=True, + save_top_k=1, + monitor="val/acc_top1", + mode="max" + ), + + LearningRateMonitor( + logging_interval="epoch" + ), + + EarlyStopping( + monitor="val/loss", + min_delta=2e-4, + patience=10, + verbose=False, + mode="min" + ) +] + +# Системные метрики +mlflow.enable_system_metrics_logging() +mlflow.set_system_metrics_sampling_interval(5) + +logger = MLFlowLogger( + experiment_name=os.getenv("EXPERIMENT_NAME"), + run_name="VGG16Custom", + tracking_uri=os.getenv("TRACKING_URI"), + log_model=True, +) + + +vgg = VGGCustom(architecture="VGG16", + num_channels=3, + num_classes=1000, + batch_normalization=False, + dropout=.5) + +# Необходимо чтобы отрпаботал LazyLinear и принял размерности! +# with torch.no_grad(): +# alexnet(torch.zeros(1, 3, 224, 224)) + +model = ClassificationLightning(model=vgg, + num_classes=1000, + optimizer_class=optim.Adam, + optimizer_kwargs=optimizer_Adam_kwargs, + scheduler_class=optim.lr_scheduler.ReduceLROnPlateau, + scheduler_kwargs=scheduler_ReduceLR_kwargs, + scheduler_monitor="val/loss") + +trainer = pl.Trainer( + max_epochs=num_epoch, + accelerator='gpu', + devices=1, + precision="16-mixed", + logger=logger, + callbacks= callbacks +) + +trainer.fit(model=model, + datamodule=imagenet_dm) \ No newline at end of file diff --git a/train/__init__.py b/train/__init__.py new file mode 100644 index 0000000..94d3f79 --- /dev/null +++ b/train/__init__.py @@ -0,0 +1,6 @@ +from .classification import ClassificationLightning + + +__all__ = [ + "ClassificationLightning" +] \ No newline at end of file diff --git a/train/classification.py b/train/classification.py new file mode 100644 index 0000000..857d9fe --- /dev/null +++ b/train/classification.py @@ -0,0 +1,165 @@ +import time + +import torch +import torch.nn as nn +import pytorch_lightning as pl + +from torch import optim +from torchmetrics import MetricCollection +from torchmetrics.classification import Accuracy + + +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 + + 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 + } +