Добавлен AlexNet

BaseCNNArhitecture/AlexNetCustom.py

@@ -3,15 +3,21 @@
 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, random_split
+from torch.utils.data import DataLoader, Dataset
 
+# Логирование данных
+from pytorch_lightning.loggers import MLFlowLogger
 # автоматическое отслеживание lr
 from pytorch_lightning.callbacks import LearningRateMonitor
 # ранние остановки
@@ -20,10 +26,11 @@ from pytorch_lightning.callbacks.early_stopping import EarlyStopping
 from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
 # %% Параметры
 
-dataset_root = "./Datasets/ImageNet"
-batch_size = 64
+dataset_root = "/media/daniel/developments/Datasets/ImageNet"
+batch_size = 128
 num_workers = 8
-num_epoch = 2
+num_epoch = 90
+learning_rate = 3e-4
 
 # %% Трансформации
 
@@ -59,62 +66,136 @@ path = kagglehub.competition_download('imagenet-object-localization-challenge',
 print("Path to competition files:", path)
 # %% Датасет
 
-full_train_dataset = datasets.ImageNet(
-    root=dataset_root,
-    split="train",
-    transform=train_transform
-)
+TransformType = Optional[Callable[..., Any]]
 
-train_size = int(0.9 * len(full_train_dataset))
-val_size = len(full_train_dataset) - train_size
+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
 
-train_dataset, val_dataset = random_split(
-    full_train_dataset,
-    [train_size, val_size],
-    generator=torch.Generator().manual_seed(42)
-)
+        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
 
 
-val_base_dataset = datasets.ImageNet(
-    root=dataset_root,
-    split="train",
-    transform=eval_transform
-)
+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"
 
-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
-)
+        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
 
@@ -163,7 +244,7 @@ class AlexNetCustom(nn.Module):
         
     def forward(self, image):
         x = self.feature(image)
-        # x = self.avgpool(x)
+        x = self.avgpool(x)
         x = torch.flatten(x, 1)
         x = self.mlp(x)
         
@@ -241,9 +322,6 @@ class ClassificationLightning(pl.LightningModule):
                     stage,
                     metrics):
         x, y = sample
-        x = x.to(torch.float32)
-        y = y.to(torch.int64)
-        
         
         logits = self(x)
         loss = self.criterion(logits, y)
@@ -327,3 +405,84 @@ class ClassificationLightning(pl.LightningModule):
             "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)