Обнавлена структура проекта

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
+        }
+