feat(train, dist_train, plots): improved training, implemented distributed training, added plotting

performance.json

@@ -0,0 +1,122 @@
+[
+    {
+        "avg_train_loss": 2.0,
+        "train_accuracy": 0.0,
+        "avg_val_loss": 4.0,
+        "val_accuracy": 0.0
+    },
+    {
+        "avg_train_loss": 1.3333333333333333,
+        "train_accuracy": 0.125,
+        "avg_val_loss": 2.0,
+        "val_accuracy": 0.1
+    },
+    {
+        "avg_train_loss": 1.0,
+        "train_accuracy": 0.2222222222222222,
+        "avg_val_loss": 1.3333333333333333,
+        "val_accuracy": 0.18181818181818182
+    },
+    {
+        "avg_train_loss": 0.8,
+        "train_accuracy": 0.3,
+        "avg_val_loss": 1.0,
+        "val_accuracy": 0.25
+    },
+    {
+        "avg_train_loss": 0.6666666666666666,
+        "train_accuracy": 0.36363636363636365,
+        "avg_val_loss": 0.8,
+        "val_accuracy": 0.3076923076923077
+    },
+    {
+        "avg_train_loss": 0.5714285714285714,
+        "train_accuracy": 0.4166666666666667,
+        "avg_val_loss": 0.6666666666666666,
+        "val_accuracy": 0.35714285714285715
+    },
+    {
+        "avg_train_loss": 0.5,
+        "train_accuracy": 0.46153846153846156,
+        "avg_val_loss": 0.5714285714285714,
+        "val_accuracy": 0.4
+    },
+    {
+        "avg_train_loss": 0.4444444444444444,
+        "train_accuracy": 0.5,
+        "avg_val_loss": 0.5,
+        "val_accuracy": 0.4375
+    },
+    {
+        "avg_train_loss": 0.4,
+        "train_accuracy": 0.5333333333333333,
+        "avg_val_loss": 0.4444444444444444,
+        "val_accuracy": 0.47058823529411764
+    },
+    {
+        "avg_train_loss": 0.36363636363636365,
+        "train_accuracy": 0.5625,
+        "avg_val_loss": 0.4,
+        "val_accuracy": 0.5
+    },
+    {
+        "avg_train_loss": 0.3333333333333333,
+        "train_accuracy": 0.5882352941176471,
+        "avg_val_loss": 0.36363636363636365,
+        "val_accuracy": 0.5263157894736842
+    },
+    {
+        "avg_train_loss": 0.3076923076923077,
+        "train_accuracy": 0.6111111111111112,
+        "avg_val_loss": 0.3333333333333333,
+        "val_accuracy": 0.55
+    },
+    {
+        "avg_train_loss": 0.2857142857142857,
+        "train_accuracy": 0.631578947368421,
+        "avg_val_loss": 0.3076923076923077,
+        "val_accuracy": 0.5714285714285714
+    },
+    {
+        "avg_train_loss": 0.26666666666666666,
+        "train_accuracy": 0.65,
+        "avg_val_loss": 0.2857142857142857,
+        "val_accuracy": 0.5909090909090909
+    },
+    {
+        "avg_train_loss": 0.25,
+        "train_accuracy": 0.6666666666666666,
+        "avg_val_loss": 0.26666666666666666,
+        "val_accuracy": 0.6086956521739131
+    },
+    {
+        "avg_train_loss": 0.23529411764705882,
+        "train_accuracy": 0.6818181818181818,
+        "avg_val_loss": 0.25,
+        "val_accuracy": 0.625
+    },
+    {
+        "avg_train_loss": 0.2222222222222222,
+        "train_accuracy": 0.6956521739130435,
+        "avg_val_loss": 0.23529411764705882,
+        "val_accuracy": 0.64
+    },
+    {
+        "avg_train_loss": 0.21052631578947367,
+        "train_accuracy": 0.7083333333333334,
+        "avg_val_loss": 0.2222222222222222,
+        "val_accuracy": 0.6538461538461539
+    },
+    {
+        "avg_train_loss": 0.2,
+        "train_accuracy": 0.72,
+        "avg_val_loss": 0.21052631578947367,
+        "val_accuracy": 0.6666666666666666
+    },
+    {
+        "avg_train_loss": 0.19047619047619047,
+        "train_accuracy": 0.7307692307692307,
+        "avg_val_loss": 0.2,
+        "val_accuracy": 0.6785714285714286
+    }
+]

plots.py

@@ -0,0 +1,42 @@
+import json
+import matplotlib.pyplot as plt
+
+with open("performance.json", "r") as f:
+    performance = json.load(f)
+
+# Extract values from the performance list
+epochs = range(1, len(performance) + 1)
+train_losses = [epoch["avg_train_loss"] for epoch in performance]
+val_losses = [epoch["avg_val_loss"] for epoch in performance]
+train_accuracies = [epoch["train_accuracy"] for epoch in performance]
+val_accuracies = [epoch["val_accuracy"] for epoch in performance]
+
+# Plot Training and Validation Loss
+plt.figure(figsize=(14, 6))
+
+# Subplot for Loss
+plt.subplot(1, 2, 1)
+plt.plot(epochs, train_losses, label="Training Loss")
+plt.plot(epochs, val_losses, label="Validation Loss")
+plt.xlabel("Epochs")
+plt.ylabel("Loss")
+plt.title("Training and Validation Loss")
+plt.legend()
+plt.xticks(epochs)
+
+# Subplot for Accuracy
+plt.subplot(1, 2, 2)
+plt.plot(epochs, train_accuracies, label="Training Accuracy")
+plt.plot(epochs, val_accuracies, label="Validation Accuracy")
+plt.xlabel("Epochs")
+plt.ylabel("Accuracy")
+plt.title("Training and Validation Accuracy")
+plt.legend()
+plt.xticks(epochs)
+
+plt.tight_layout()
+
+# Save the plot as an image file
+plt.savefig("performance_plot.png", dpi=300)
+
+plt.show()

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+pillow
+tqdm
+matplotlib

scene_classification.py → train.py

@@ -1,5 +1,7 @@
+#!/usr/bin/env python3
 import os
 import csv
+import json
 from tqdm import tqdm
 import torch
 import argparse
@@ -134,6 +136,15 @@ def train(model, train_loader, val_loader, optimizer, criterion, device,
     # Place the model on device
     model = model.to(device)
 
+    # Define early stopping parameters
+    patience = 3  # 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
+
+    # Performance tracking
+    performance = []
+
     for epoch in range(num_epochs):
         model.train()  # Set model to training mode
 
@@ -176,15 +187,45 @@ def train(model, train_loader, val_loader, optimizer, criterion, device,
                 pbar.set_postfix(loss=loss.item())
 
             # Calculate average loss and accuracy
-            avg_loss = running_loss / len(train_loader)
-            accuracy = correct_predictions / total_samples
-
-            avg_val_loss, val_accuracy  = evaluate(model, val_loader, criterion, device)
+            avg_train_loss = running_loss / len(train_loader)
+            train_accuracy = correct_predictions / total_samples
+            avg_val_loss, val_accuracy = evaluate(model, val_loader, criterion, device)
+
+            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_loss:.4f}, Accuracy: {accuracy:.4f} "
+                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
+
+    # Save the performance list to a JSON file
+    with open("performance.json", "w") as f:
+        json.dump(performance, f, indent=4)
+    torch.save(best_model_state, 'model.ckpt')
+
 
 def test(model, test_loader, device):
     """
@@ -250,7 +291,7 @@ def main(args):
     # Create the dataloaders
 
     # Define the batch size and number of workers
-    batch_size = 64
+    batch_size = int(args.batch_size)
     num_workers = 2
 
     # Create DataLoader for training and validation sets
@@ -263,7 +304,7 @@ def main(args):
                             num_workers=num_workers,
                             shuffle=False)
 
-    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')  # TODO: check cuda
+    device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else 'cpu')  # TODO: check cuda
 
     model = MyModel(num_classes=len(miniplaces_train.label_dict))
 
@@ -279,10 +320,7 @@ def main(args):
 
     if not args.test:
         train(model, train_loader, val_loader, optimizer, criterion,
-              device, num_epochs=25)
-
-        torch.save({'model_state_dict': model.state_dict(),
-                    'optimizer_state_dict': optimizer.state_dict()}, 'model.ckpt')
+              device, num_epochs=int(args.epochs))
 
     else:
         miniplaces_test = MiniPlaces(data_root,
@@ -301,6 +339,9 @@ def main(args):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('--test', action='store_true')
-    parser.add_argument('--checkpoint', default='model.ckpt')
+    parser.add_argument('--checkpoint')
+    parser.add_argument('--gpu', default=0)
+    parser.add_argument('--epochs', default=10)
+    parser.add_argument('--batch_size', default=32)
     args = parser.parse_args()
     main(args)

train_dist.py

@@ -0,0 +1,364 @@
+#!/usr/bin/env python3
+import os
+import csv
+import json
+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.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 the distributed training environment by destroying the process group.
+    """
+    dist.destroy_process_group()
+
+
+class MiniPlaces(Dataset):
+    # Your existing MiniPlaces class implementation remains the same
+    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 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_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()
+            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_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_samples, op=dist.ReduceOp.SUM)
+
+        avg_loss = (total_loss / world_size).item() / len(test_loader)
+        accuracy = (num_correct / num_samples).item()
+
+    return avg_loss, 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.
+    """
+    setup(rank, world_size, args.port)
+    device = torch.device(f'cuda:{rank}')
+
+    # Define early stopping parameters
+    patience = 3  # 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
+
+    # Data loading and preprocessing
+    image_net_mean = torch.Tensor([0.485, 0.456, 0.406])
+    image_net_std = torch.Tensor([0.229, 0.224, 0.225])
+    data_transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Resize((128, 128)),
+        transforms.Normalize(image_net_mean, image_net_std),
+    ])
+
+    # Create datasets
+    data_root = 'data'
+    miniplaces_train = MiniPlaces(data_root, split='train', transform=data_transform)
+    miniplaces_val = MiniPlaces(data_root, split='val', transform=data_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))
+    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:
+        map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
+        checkpoint = torch.load(args.checkpoint, map_location=map_location)
+        model.module.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+
+    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 = evaluate(model, val_loader, criterion, device)
+
+            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
+        miniplaces_test = MiniPlaces(data_root, split='test', transform=data_transform)
+        test_loader = DataLoader(miniplaces_test, batch_size=args.batch_size, num_workers=2, shuffle=False)
+        checkpoint = torch.load(args.checkpoint, map_location=device)
+        model.module.load_state_dict(checkpoint['model_state_dict'])
+        preds = test(model, test_loader, device)
+        if rank == 0:  # Only write predictions on rank 0
+            write_predictions(preds, 'predictions.csv')
+
+    cleanup()
+
+
+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.
+    """
+    # Get number of available GPUs
+    world_size = torch.cuda.device_count()
+    mp.spawn(train_worker,
+             args=(world_size, args),
+             nprocs=world_size,
+             join=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--test', action='store_true')
+    parser.add_argument('--checkpoint')
+    parser.add_argument('--epochs', type=int, default=10)
+    parser.add_argument('--batch_size', type=int, default=64)
+    parser.add_argument('--port', type=int, default=4224)
+    args = parser.parse_args()
+    main(args)