# %% Импорт библиотек 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 }