first commit

2026-09-14 11:05:56 +10:00
commit 41c9010db2
8 changed files with 1009 additions and 0 deletions
@@ -0,0 +1,229 @@
+# %% Импорт библиотек
+
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+
+from torch.utils.data import DataLoader
+from torchmetrics import MetricCollection
+from torchvision import datasets, transforms
+from torchmetrics.classification import Accuracy, Precision, Recall, F1Score
+
+# Автоматическое логирование результатов
+from pytorch_lightning.loggers import WandbLogger, TensorBoardLogger, CSVLogger
+# сохранение весов
+from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
+# автоматическое отслеживание lr
+from pytorch_lightning.callbacks import LearningRateMonitor
+# ранние остановки
+from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+
+# %% Гиперпараметры
+
+num_epoch = 8
+batch_size = 64
+
+# %% Создание модели PyTorch
+
+class MyTorchCNNModel(nn.Module):
+    def __init__(self, 
+                 output_size: int = 10):
+        super().__init__()
+        
+        self.output_size = output_size
+        
+        self.feature = nn.Sequential(
+            nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2)
+        )
+        
+        self.head = nn.Conv2d(256, self.output_size, kernel_size=1)
+        self.pool = nn.AdaptiveAvgPool2d((1, 1))
+        
+        
+    def forward(self, image):
+        x = self.feature(image)
+        x = self.head(x)
+        x = self.pool(x)
+        
+        return x.flatten(1)
+
+# %% Создание модели PyTorch Lightning
+
+class MyNNExample(pl.LightningModule):
+    def __init__(self,
+                 model: nn.Module,
+                 eta=3e-4):
+        super().__init__()
+        
+        self.save_hyperparameters(ignore=["model"])
+        # device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = model
+        
+        self.criterion = nn.CrossEntropyLoss()
+        self.train_metrics = MetricCollection({
+            "accuracy": Accuracy(task="multiclass", 
+                                 num_classes=model.output_size),
+            "precision": Precision(task="multiclass", 
+                                   num_classes=model.output_size, 
+                                   average="macro"),    
+            "recall": Recall(task="multiclass", 
+                             num_classes=model.output_size, 
+                             average="macro"),
+            "f1": F1Score(task="multiclass", 
+                          num_classes=model.output_size, 
+                          average="macro")    
+        })
+        
+        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.forward(x.to(torch.float32))
+        preds = torch.argmax(logits, 1)
+        loss = self.criterion(logits, y.to(torch.int64))
+        
+        metrics.update(preds, 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 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)
+    
+    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 = torch.optim.Adam(self.parameters(), lr=self.hparams.eta)
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, patience=5, mode="min"
+        )
+        
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "monitor": "val_loss",
+                "interval": "epoch",
+            }
+        }
+
+# %% Загрузка набора данных
+
+image_transforms = transforms.Compose([
+    transforms.ToTensor()
+])
+
+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)
+
+# %% Запуск обучения
+
+model = MyNNExample(model=MyTorchCNNModel())
+callbacks = [
+    ModelCheckpoint(dirpath="models",
+                    filename="{epoch}_{val_accuracy:.4f}_{val_loss:.4f}",
+                    auto_insert_metric_name=False,
+                    save_last=True,
+                    save_top_k=1,
+                    monitor="val_loss",
+                    mode="min"),
+    LearningRateMonitor(logging_interval="step"),
+    EarlyStopping(monitor="val_loss",
+                  min_delta=2e-4,
+                  patience=10,
+                  verbose=False,
+                  mode="min")
+]
+
+trainer = pl.Trainer(
+    max_epochs=num_epoch,
+    accelerator="gpu",
+    devices=1,
+    logger=CSVLogger("./train_logs", name="mnist_csv", 
+                     version=None, flush_logs_every_n_steps=100),
+    callbacks=callbacks
+)
+
+trainer.fit(model=model, 
+            train_dataloaders=train_loader,
+            val_dataloaders=test_loader)
+        
+# %% Инференс модели
+        
+checkpoint_path = "./models/6_0.9460_0.1643.ckpt"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = MyNNExample.load_from_checkpoint(checkpoint_path,
+                                         model=MyTorchCNNModel(output_size=10))
+model = model.to(device)
+model.eval()
+
+image, true_label = test_dataset[42]
+image_batch = image.unsqueeze(0).to(device)
+
+with torch.no_grad():
+    logits = model(image_batch)
+    probs = torch.softmax(logits, dim=1)
+    pred_label = logits.argmax(dim=1).item()
+    
+print(f"Предсказание: {pred_label}")
+print(f"Истинная метка: {true_label}")
+print(f"Уверенность: {probs[0, pred_label].item():.2%}")