DistributedNetwork.py

import os
import csv
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.cuda.amp import autocast, GradScaler
from torchvision import datasets
from Checkpoint_utils import save_checkpoint

from CNN_Architecture import alexNet, save_model

torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision("high")

# =========================
# EDIT THESE LOCAL SETTINGS
# =========================
TRAIN_DIR = "./train"                 # ImageFolder root
VALID_DIR = "./valid"                 # ImageFolder root

LOG_CSV_PATH = "training_log.csv"
OVERWRITE_LOG = True

EPOCHS = 90
BATCH_SIZE = 64
LR = 1e-2
WEIGHT_DECAY = 5e-4

NUM_WORKERS = 1
PIN_MEMORY = True

TRAIN_TRANSFORM_PATH = "train_transform.pth"
VAL_TRANSFORM_PATH = "val_transform.pth"

PLATEAU_PATIENCE = 4
MIN_LR = 1e-4
# =========================


class Trainer:
    def __init__(self):
        # ---- Distributed setup ----
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        local_rank = int(os.environ["LOCAL_RANK"])

        torch.cuda.set_device(local_rank)
        dist.init_process_group(backend="nccl")

        self.rank = rank
        self.world_size = world_size
        self.local_rank = local_rank
        self.device = torch.device("cuda", local_rank)
        self.checkpoint_dir = "Checkpoints"

        # ---- Load transforms ----
        self.train_transforms = torch.load(TRAIN_TRANSFORM_PATH, map_location="cpu")
        self.valid_transforms = torch.load(VAL_TRANSFORM_PATH, map_location="cpu")

        # ---- Datasets ----
        train_dataset = datasets.ImageFolder(
            root=TRAIN_DIR,
            transform=self.train_transforms
        )

        valid_dataset = datasets.ImageFolder(
            root=VALID_DIR,
            transform=self.valid_transforms
        )

        # Store class mapping so inference can map indices -> class names
        self.class_to_idx = train_dataset.class_to_idx

        # ---- Samplers ----
        self.trainSampler = DistributedSampler(
            train_dataset, num_replicas=world_size, rank=rank, shuffle=True
        )
        self.validSampler = DistributedSampler(
            valid_dataset, num_replicas=world_size, rank=rank, shuffle=False
        )

        # ---- DataLoaders ----
        self.trainLoader = DataLoader(
            train_dataset,
            batch_size=BATCH_SIZE,
            sampler=self.trainSampler,
            num_workers=4,
            pin_memory=PIN_MEMORY,
            persistent_workers=True,
            prefetch_factor=4,
        )

        self.validLoader = DataLoader(
            valid_dataset,
            batch_size=BATCH_SIZE,
            sampler=self.validSampler,
            num_workers=4,
            pin_memory=PIN_MEMORY,
            persistent_workers=True,
            prefetch_factor=4,
        )

        # ---- Model / optimizer / loss ----
        model = alexNet().to(self.device)
        self.model = DDP(model, device_ids=[local_rank],static_graph=True)

        self.optimizer = optim.Adam(
            self.model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY, betas=(0.9, 0.999)
        )

        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer,
            mode="min",       # minimize validation loss
            factor=0.1,       # divide LR by 10
            patience=PLATEAU_PATIENCE,       # how many epochs to wait with no improvement
            threshold=1e-4,   # minimum change to qualify as improvement
            cooldown=0,       # epochs to wait after LR reduction
            min_lr=MIN_LR,
        )

        # Uses logits directly; internally applies log_softmax + NLLLoss
        self.train_criterion = nn.CrossEntropyLoss(reduction="none")

        # For paper-faithful 10-crop validation we’ll compute loss from averaged probs
        # (negative log likelihood on averaged probability of true class)
        self.eps = 1e-12

        self._init_csv_logger()

    def _init_csv_logger(self):
        if self.rank != 0:
            return
        mode = "w" if OVERWRITE_LOG else "a"
        file_exists = os.path.exists(LOG_CSV_PATH)
        if mode == "w" or not file_exists:
            with open(LOG_CSV_PATH, mode, newline="") as f:
                writer = csv.writer(f)
                writer.writerow(["epoch", "lr", "train_loss", "val_loss", "val_acc"])

    def validation(self):
        self.model.eval()

        local_loss_sum = 0.0
        local_sample_count = 0
        local_correct = 0

        amp_enabled = (self.device.type == "cuda")

        with torch.no_grad():
            for images, labels in self.validLoader:
                images = images.to(self.device, non_blocking=True)
                labels = labels.to(self.device, non_blocking=True)

                with autocast("cuda", dtype=torch.float16): # type: ignore
                    if images.dim() == 5:
                        # TenCrop: images shape (B, 10, 3, 224, 224)
                        B, NC, C, H, W = images.shape
                        flat = images.view(B * NC, C, H, W)  # (B*10, 3, 224, 224)

                        logits = self.model(flat)            # (B*10, K)

                        # Paper-style: average softmax probabilities across crops
                        probs = torch.softmax(logits, dim=1)              # (B*10, K)
                        probs = probs.view(B, NC, -1).mean(dim=1)         # (B, K)

                        # Accuracy
                        preds = probs.argmax(dim=1)
                        local_correct += (preds == labels).sum().item()

                        # Loss: -log(prob_true)
                        prob_true = probs.gather(1, labels.view(-1, 1)).squeeze(1)  # (B,)
                        losses = -torch.log(prob_true.clamp_min(self.eps))         # (B,)

                    else:
                        # Non-TenCrop fallback
                        logits = self.model(images)  # (B, K)
                        losses = self.train_criterion(logits, labels)
                        preds = logits.argmax(dim=1)
                        local_correct += (preds == labels).sum().item()

                local_loss_sum += losses.sum().item()
                local_sample_count += labels.size(0)
                local_correct += (preds == labels).sum().item()

        # ---- All-reduce across ranks ----
        loss_tensor = torch.tensor(local_loss_sum, device=self.device, dtype=torch.float64)
        count_tensor = torch.tensor(local_sample_count, device=self.device, dtype=torch.float64)
        correct_tensor = torch.tensor(local_correct, device=self.device, dtype=torch.float64)

        dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
        dist.all_reduce(count_tensor, op=dist.ReduceOp.SUM)
        dist.all_reduce(correct_tensor, op=dist.ReduceOp.SUM)

        avg_loss = (loss_tensor / count_tensor).item()
        avg_acc = (correct_tensor / count_tensor).item()
        return avg_loss, avg_acc

    def train(self):
        is_cuda = (self.device.type == "cuda")
        scaler = GradScaler("cuda", enabled=is_cuda) # type: ignore

        best_val_loss = float("inf")
        ckpt_dir = getattr(self, "checkpoint_dir", "Checkpoints")
        os.makedirs(ckpt_dir, exist_ok=True)

        for epoch in range(EPOCHS):
            self.model.train()
            self.trainSampler.set_epoch(epoch)

            local_loss_sum = 0.0
            local_sample_count = 0

            for imgs, targets in self.trainLoader:
                imgs = imgs.to(self.device, non_blocking=True)
                targets = targets.to(self.device, non_blocking=True)

                self.optimizer.zero_grad(set_to_none=True)

                with autocast("cuda", dtype=torch.float16): # type: ignore
                    logits = self.model(imgs)                       # (B, K)
                    losses = self.train_criterion(logits, targets)  # (B,)
                    loss = losses.mean()

                scaler.scale(loss).backward()
                scaler.step(self.optimizer)
                scaler.update()

                local_loss_sum += losses.sum().item()
                local_sample_count += targets.size(0)

            train_loss_tensor = torch.tensor(local_loss_sum, device=self.device, dtype=torch.float64)
            train_count_tensor = torch.tensor(local_sample_count, device=self.device, dtype=torch.float64)

            dist.all_reduce(train_loss_tensor, op=dist.ReduceOp.SUM)
            dist.all_reduce(train_count_tensor, op=dist.ReduceOp.SUM)

            train_loss_global = (train_loss_tensor / train_count_tensor).item()

            valid_loss, valid_acc = self.validation()
            self.scheduler.step(valid_loss)

            if self.rank == 0:
                current_lr = self.optimizer.param_groups[0]["lr"]

                # ---- Safe CSV writing ----
                try:
                    with open(LOG_CSV_PATH, "a", newline="") as f:
                        writer = csv.writer(f)
                        writer.writerow([epoch, current_lr, train_loss_global, valid_loss, valid_acc])
                except Exception as e:
                    print(f"[Warning] CSV logging failed: {e} | Epoch = {epoch}")

                model_to_save = self.model.module if hasattr(self.model, "module") else self.model

                class_to_idx = getattr(self, "class_to_idx", None)

                # save per-epoch
                save_checkpoint(
                    os.path.join(ckpt_dir, f"epoch_{epoch:03d}.pth"),
                    epoch=epoch,
                    model=model_to_save,
                    optimizer=self.optimizer,
                    scheduler=self.scheduler,
                    scaler=scaler,
                    class_to_idx=class_to_idx,
                    extra={
                        "train_loss_global": train_loss_global,
                        "valid_loss_global": float(valid_loss),
                        "valid_acc_global": float(valid_acc),
                    },
                )

                # best model by val loss
                if valid_loss < best_val_loss:
                    best_val_loss = valid_loss
                    save_checkpoint(
                        os.path.join(ckpt_dir, "best.pth"),
                        epoch=epoch,
                        model=model_to_save,
                        optimizer=self.optimizer,
                        scheduler=self.scheduler,
                        scaler=scaler,
                        class_to_idx=class_to_idx,
                        extra={
                            "train_loss_global": train_loss_global,
                            "valid_loss_global": float(valid_loss),
                            "valid_acc_global": float(valid_acc),
                            "best_val_loss": float(best_val_loss),
                        },
                    )

                current_lr = self.optimizer.param_groups[0]["lr"]
                print(
                    f"Epoch {epoch}/{EPOCHS}  "
                    f"LR={current_lr:.6g}  "
                    f"TrainLoss(global)={train_loss_global:.4f}  "
                    f"ValLoss(global)={valid_loss:.4f}  "
                    f"ValAcc(global)={valid_acc:.4f}"
                )

        dist.destroy_process_group()


def main():
    Trainer().train()


if __name__ == "__main__":
    main()