train_dist.py

#!/usr/bin/env python3
import os
import csv
import json
import warnings
from tqdm import tqdm
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data.distributed import DistributedSampler
import argparse
from PIL import Image
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
from model import MyModel


def setup(rank, world_size, port):
    """
    Initialize the distributed training environment.

    Args:
        rank (int): The rank of the current process.
        world_size (int): The total number of processes (GPUs).
        port (int): The port number for communication.
    """
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = str(port)
    dist.init_process_group("nccl", rank=rank, world_size=world_size)


def cleanup():
    """
    Clean up distributed training environment
    """
    if dist.is_initialized():
        dist.barrier()  # Synchronize all processes before destroying process group
        dist.destroy_process_group()
        torch.cuda.synchronize()


class MiniPlaces(Dataset):
    def __init__(self, root_dir, split, transform=None, label_dict=None):
        """
        Initialize the MiniPlaces dataset with the root directory for the images,
        the split (train/val/test), an optional data transformation,
        and an optional label dictionary.

        Args:
            root_dir (str): Root directory for the MiniPlaces images.
            split (str): Split to use ('train', 'val', or 'test').
            transform (callable, optional): Optional data transformation to apply to the images.
            label_dict (dict, optional): Optional dictionary mapping integer labels to class names.
        """
        assert split in ['train', 'val', 'test']
        self.root_dir = root_dir
        self.split = split
        self.transform = transform
        self.filenames = []
        self.labels = []
        self.label_dict = label_dict if label_dict is not None else {}

        with open(os.path.join(self.root_dir, self.split + '.txt')) as r:
            lines = r.readlines()
            for line in lines:
                line = line.split()
                self.filenames.append(line[0])
                if split == 'test':
                    label = line[0]
                else:
                    label = int(line[1])
                self.labels.append(label)
                if split == 'train':
                    text_label = line[0].split('/')[2]
                    self.label_dict[label] = text_label

    def __len__(self):
        """
        Return the number of images in the dataset.

        Returns:
            int: Number of images in the dataset.
        """
        return len(self.labels)

    def __getitem__(self, idx):
        """
        Return a single image and its corresponding label when given an index.

        Args:
            idx (int): Index of the image to retrieve.

        Returns:
            tuple: Tuple containing the image and its label.
        """
        if self.transform is not None:
            image = self.transform(
                Image.open(os.path.join(self.root_dir, "images", self.filenames[idx])))
        else:
            image = Image.open(os.path.join(self.root_dir, "images", self.filenames[idx]))
        label = self.labels[idx]
        return image, label


def create_train_transform():
    """
    Create training data transformation with augmentation
    """
    image_net_mean = torch.Tensor([0.485, 0.456, 0.406])
    image_net_std = torch.Tensor([0.229, 0.224, 0.225])

    return transforms.Compose([
        transforms.RandomResizedCrop(128, scale=(0.8, 1.0)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ColorJitter(
            brightness=0.4,
            contrast=0.4,
            saturation=0.4,
            hue=0.1
        ),
        transforms.RandomAffine(
            degrees=15,  # rotation
            translate=(0.1, 0.1),  # horizontal/vertical translation
            scale=(0.9, 1.1),  # scale
        ),
        transforms.ToTensor(),
        transforms.Resize((128, 128)),
        transforms.Normalize(image_net_mean, image_net_std)
    ])


def create_val_transform():
    """
    Create validation/test data transformation without augmentation
    """
    image_net_mean = torch.Tensor([0.485, 0.456, 0.406])
    image_net_std = torch.Tensor([0.229, 0.224, 0.225])

    return transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((128, 128)),
        transforms.Normalize(image_net_mean, image_net_std)
    ])


def evaluate(model, test_loader, criterion, device):
    """
    Evaluate the CNN classifier on the validation set.

    Args:
        model (CNN): CNN classifier to evaluate.
        test_loader (torch.utils.data.DataLoader): Data loader for the test set.
        criterion (callable): Loss function to use for evaluation.
        device (torch.device): Device to use for evaluation.

    Returns:
        float: Average loss on the test set.
        float: Accuracy on the test set.
    """
    model.eval()

    with torch.no_grad():
        total_loss = 0.0
        num_correct = 0
        num_correct_top5 = 0
        num_samples = 0

        for inputs, labels in test_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)

            logits = model(inputs)
            loss = criterion(logits, labels)
            total_loss += loss.item()

            _, predictions = torch.max(logits, dim=1)
            num_correct += (predictions == labels).sum().item()

            _, top5_predictions = torch.topk(logits, 5, dim=1)
            num_correct_top5 += (top5_predictions == labels.unsqueeze(1)).any(dim=1).sum().item()

            num_samples += len(inputs)

        # Gather metrics from all processes
        world_size = dist.get_world_size()
        total_loss = torch.tensor(total_loss).to(device)
        num_correct = torch.tensor(num_correct).to(device)
        num_correct_top5 = torch.tensor(num_correct_top5).to(device)
        num_samples = torch.tensor(num_samples).to(device)

        dist.all_reduce(total_loss, op=dist.ReduceOp.SUM)
        dist.all_reduce(num_correct, op=dist.ReduceOp.SUM)
        dist.all_reduce(num_correct_top5, op=dist.ReduceOp.SUM)
        dist.all_reduce(num_samples, op=dist.ReduceOp.SUM)

        avg_loss = (total_loss / world_size).item() / len(test_loader)
        accuracy = (num_correct / num_samples).item()
        top5_accuracy = (num_correct_top5 / num_samples).item()

    return avg_loss, accuracy, top5_accuracy


def train_worker(rank, world_size, args):
    """
    Train the model in a distributed setup.

    Args:
        rank (int): The rank of the current process.
        world_size (int): The total number of processes (GPUs).
        args (argparse.Namespace): Command-line arguments.
    """
    try:
        warnings.filterwarnings("ignore")
        setup(rank, world_size, args.port)
        device = torch.device(f'cuda:{rank}')

        # Define early stopping parameters
        patience = 10  # Number of epochs to wait for improvement
        best_val_accuracy = 0.0  # Best validation accuracy so far
        epochs_without_improvement = 0  # Counter for epochs without improvement
        best_model_state = None  # To store the state of the best model

        last_lr = 0

        # Separate transforms for training and validation
        train_transform = create_train_transform()
        val_transform = create_val_transform()

        # Create datasets
        data_root = 'data'
        miniplaces_train = MiniPlaces(data_root, split='train', transform=train_transform)
        miniplaces_val = MiniPlaces(data_root, split='val', transform=val_transform,
                                    label_dict=miniplaces_train.label_dict)

        # Create distributed samplers
        train_sampler = DistributedSampler(miniplaces_train, num_replicas=world_size, rank=rank)
        val_sampler = DistributedSampler(miniplaces_val, num_replicas=world_size, rank=rank)

        # Create dataloaders
        train_loader = DataLoader(miniplaces_train, batch_size=args.batch_size,
                                  num_workers=2, sampler=train_sampler,
                                  pin_memory=True)
        val_loader = DataLoader(miniplaces_val, batch_size=args.batch_size,
                                num_workers=2, sampler=val_sampler,
                                pin_memory=True)

        # Create model and move to GPU
        model = MyModel(num_classes=len(miniplaces_train.label_dict), dropout_rate=0.2)
        model = model.to(device)
        model = DDP(model, device_ids=[rank])

        optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9,
                                    dampening=0, weight_decay=1e-4, nesterov=True)
        criterion = torch.nn.CrossEntropyLoss(reduction='mean', label_smoothing=0.1)

        if args.checkpoint or args.test:
            map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
            checkpoint = torch.load((args.checkpoint if args.checkpoint else 'model.ckpt'), map_location=map_location)
            model.module.load_state_dict(checkpoint['model_state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

        # Initialize the ReduceLROnPlateau scheduler
        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=4)

        if not args.test:
            # Training loop
            performance = []
            for epoch in range(args.epochs):
                model.train()
                train_sampler.set_epoch(epoch)  # Important for proper shuffling

                running_loss = 0.0
                correct_predictions = 0
                total_samples = 0

                if rank == 0:  # Only show progress bar on rank 0
                    pbar = tqdm(total=len(train_loader),
                                desc=f'Epoch {epoch + 1}/{args.epochs}',
                                position=0, leave=True)

                for inputs, labels in train_loader:
                    inputs = inputs.to(device)
                    labels = labels.to(device)

                    optimizer.zero_grad()
                    logits = model(inputs)
                    loss = criterion(logits, labels)
                    loss.backward()
                    optimizer.step()

                    running_loss += loss.item()
                    _, predicted = logits.max(1)
                    correct_predictions += (predicted == labels).sum().item()
                    total_samples += labels.size(0)

                    if rank == 0:
                        pbar.update(1)
                        pbar.set_postfix(loss=loss.item())

                if rank == 0:
                    pbar.close()

                # Evaluate and log metrics
                avg_train_loss = running_loss / len(train_loader)
                train_accuracy = correct_predictions / total_samples
                avg_val_loss, val_accuracy, val_top5_accuracy = evaluate(model, val_loader, criterion, device)

                # Step the scheduler with the validation loss
                scheduler.step(avg_val_loss)
                if scheduler.get_last_lr()[0] != last_lr:
                    last_lr = scheduler.get_last_lr()[0]
                    if epoch != 0:
                        print(f"New learning rate: {scheduler.get_last_lr()[0]}")

                if rank == 0:  # Only save metrics on rank 0
                    performance.append({
                        "avg_train_loss": avg_train_loss,
                        "train_accuracy": train_accuracy,
                        "avg_val_loss": avg_val_loss,
                        "val_accuracy": val_accuracy
                    })
                    print(
                        f"Train Loss: {avg_train_loss:.4f}, Accuracy: {train_accuracy:.4f} "
                        f"Validation Loss: {avg_val_loss:.4f}, Validation Accuracy: {val_accuracy:.4f}"
                    )

                    # Check for early stopping
                    if val_accuracy > best_val_accuracy:
                        best_val_accuracy = val_accuracy
                        epochs_without_improvement = 0  # Reset counter if there's an improvement

                        # Save the model checkpoint for the best model
                        best_model_state = {
                            'model_state_dict': model.module.state_dict(),
                            'optimizer_state_dict': optimizer.state_dict(),
                            'epoch': epoch,
                        }
                    else:
                        epochs_without_improvement += 1

                    # Early stopping condition
                    if epochs_without_improvement >= patience:
                        print(f"Early stopping at epoch {epoch + 1}.")
                        break  # Stop training if no improvement for 'patience' epochs

            if rank == 0:  # Save performance and the best model checkpoint only on rank 0
                with open("performance.json", "w") as f:
                    json.dump(performance, f, indent=4)
                torch.save(best_model_state, 'model.ckpt')

        else:  # Testing mode
            avg_val_loss, val_accuracy, val_top5_accuracy = evaluate(model, val_loader, criterion, device)
            if rank == 0:
                print(f"\nValidation Loss: {avg_val_loss:.4f}\n"
                      f"Validation Accuracy: {val_accuracy:.4f}\n"
                      f"Validation Top-5 Accuracy: {val_top5_accuracy:.4f}\n")

            miniplaces_test = MiniPlaces(data_root, split='test', transform=val_transform)
            test_loader = DataLoader(miniplaces_test, batch_size=args.batch_size, num_workers=2, shuffle=False)

            preds = test(model, test_loader, device)
            if rank == 0:  # Only write predictions on rank 0
                write_predictions(preds, 'predictions.csv')
                print("Predictions saved to predictions.csv\n")

    finally:
        cleanup()
        # Explicit synchronization before exiting
        torch.cuda.synchronize()
        if dist.is_initialized():
            dist.barrier()


def test(model, test_loader, device):
    """
    Test the model on a dataset and return predictions.

    Args:
        model (torch.nn.Module): The model to test.
        test_loader (DataLoader): The DataLoader for the test dataset.
        device (torch.device): The device to run the test on.

    Returns:
        list: A list of (label, prediction) tuples for each image.
    """
    model.eval()
    with torch.no_grad():
        all_preds = []
        for inputs, labels in test_loader:
            inputs = inputs.to(device)
            logits = model(inputs)
            _, predictions = torch.max(logits, dim=1)
            preds = list(zip(labels, predictions.tolist()))
            all_preds.extend(preds)
        return all_preds


def write_predictions(preds, filename):
    """
    Write model predictions to a CSV file.

    Args:
        preds (list): A list of (label, prediction) tuples.
        filename (str): The name of the CSV file to save predictions to.
    """
    with open(filename, 'w') as f:
        writer = csv.writer(f, delimiter=',')
        for im, pred in preds:
            writer.writerow((im, pred))


def main(args):
    """
    Main function to start the training process using multiple GPUs.

    Args:
        args (argparse.Namespace): Command-line arguments.
    """
    world_size = torch.cuda.device_count()
    try:
        mp.spawn(train_worker,
                 args=(world_size, args),
                 nprocs=world_size,
                 join=True)
    finally:
        # Force cleanup of any remaining CUDA resources
        torch.cuda.empty_cache()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--test', action='store_true')
    parser.add_argument('--checkpoint')
    parser.add_argument('--epochs', type=int, default=200)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--port', type=int, default=4224)
    args = parser.parse_args()
    main(args)