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

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)