feat(data): implemented data augmentation

model.py

@@ -34,9 +34,9 @@ class MyModel(nn.Module):
         self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
 
         # Residual blocks
-        self.layer1 = self._resnet_layers(64, 128, num_blocks=2)  # 2 residual blocks
-        self.layer2 = self._resnet_layers(128, 256, num_blocks=2)  # 2 residual blocks
-        self.layer3 = self._resnet_layers(256, 512, num_blocks=2)  # 2 residual blocks
+        self.block1 = self._resnet_layers(64, 128, num_blocks=3)  # 3 residual blocks
+        self.block2 = self._resnet_layers(128, 256, num_blocks=3)  # 3 residual blocks
+        self.block3 = self._resnet_layers(256, 512, num_blocks=3)  # 3 residual blocks
 
         # Global average pooling
         self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
@@ -47,9 +47,9 @@ class MyModel(nn.Module):
             self.bn1,
             nn.ReLU(),
             self.pool1,
-            self.layer1,
-            self.layer2,
-            self.layer3,
+            self.block1,
+            self.block2,
+            self.block3,
             self.global_avg_pool
         )
 

performance.json

@@ -1,122 +1,98 @@
 [
     {
-        "avg_train_loss": 2.0,
-        "train_accuracy": 0.0,
-        "avg_val_loss": 4.0,
-        "val_accuracy": 0.0
+        "avg_train_loss": 3.6829103430493553,
+        "train_accuracy": 0.1709,
+        "avg_val_loss": 3.5155134261793393,
+        "val_accuracy": 0.21660000085830688
     },
     {
-        "avg_train_loss": 1.3333333333333333,
-        "train_accuracy": 0.125,
-        "avg_val_loss": 2.0,
-        "val_accuracy": 0.1
+        "avg_train_loss": 3.1779507774614175,
+        "train_accuracy": 0.28584,
+        "avg_val_loss": 3.3872365769307327,
+        "val_accuracy": 0.26499998569488525
     },
     {
-        "avg_train_loss": 1.0,
-        "train_accuracy": 0.2222222222222222,
-        "avg_val_loss": 1.3333333333333333,
-        "val_accuracy": 0.18181818181818182
+        "avg_train_loss": 2.948077251571001,
+        "train_accuracy": 0.3488,
+        "avg_val_loss": 2.960327925955414,
+        "val_accuracy": 0.35409998893737793
     },
     {
-        "avg_train_loss": 0.8,
-        "train_accuracy": 0.3,
-        "avg_val_loss": 1.0,
-        "val_accuracy": 0.25
+        "avg_train_loss": 2.7825030597142506,
+        "train_accuracy": 0.39572,
+        "avg_val_loss": 2.9160548896546574,
+        "val_accuracy": 0.3675999939441681
     },
     {
-        "avg_train_loss": 0.6666666666666666,
-        "train_accuracy": 0.36363636363636365,
-        "avg_val_loss": 0.8,
-        "val_accuracy": 0.3076923076923077
+        "avg_train_loss": 2.6581287719619175,
+        "train_accuracy": 0.43032,
+        "avg_val_loss": 2.8124696768013533,
+        "val_accuracy": 0.39629998803138733
     },
     {
-        "avg_train_loss": 0.5714285714285714,
-        "train_accuracy": 0.4166666666666667,
-        "avg_val_loss": 0.6666666666666666,
-        "val_accuracy": 0.35714285714285715
+        "avg_train_loss": 2.536289040659455,
+        "train_accuracy": 0.46174,
+        "avg_val_loss": 2.7144464383459397,
+        "val_accuracy": 0.42500001192092896
     },
     {
-        "avg_train_loss": 0.5,
-        "train_accuracy": 0.46153846153846156,
-        "avg_val_loss": 0.5714285714285714,
-        "val_accuracy": 0.4
+        "avg_train_loss": 2.440945129400633,
+        "train_accuracy": 0.49412,
+        "avg_val_loss": 2.745724817749801,
+        "val_accuracy": 0.4189999997615814
     },
     {
-        "avg_train_loss": 0.4444444444444444,
-        "train_accuracy": 0.5,
-        "avg_val_loss": 0.5,
-        "val_accuracy": 0.4375
+        "avg_train_loss": 2.3424960819483567,
+        "train_accuracy": 0.52302,
+        "avg_val_loss": 2.744392152045183,
+        "val_accuracy": 0.4237000048160553
     },
     {
-        "avg_train_loss": 0.4,
-        "train_accuracy": 0.5333333333333333,
-        "avg_val_loss": 0.4444444444444444,
-        "val_accuracy": 0.47058823529411764
+        "avg_train_loss": 2.245347209489277,
+        "train_accuracy": 0.5516,
+        "avg_val_loss": 2.7382394584121217,
+        "val_accuracy": 0.43230000138282776
     },
     {
-        "avg_train_loss": 0.36363636363636365,
-        "train_accuracy": 0.5625,
-        "avg_val_loss": 0.4,
-        "val_accuracy": 0.5
+        "avg_train_loss": 2.155752474042901,
+        "train_accuracy": 0.57972,
+        "avg_val_loss": 2.7085890071407244,
+        "val_accuracy": 0.43689998984336853
     },
     {
-        "avg_train_loss": 0.3333333333333333,
-        "train_accuracy": 0.5882352941176471,
-        "avg_val_loss": 0.36363636363636365,
-        "val_accuracy": 0.5263157894736842
+        "avg_train_loss": 2.0571537492218797,
+        "train_accuracy": 0.6087,
+        "avg_val_loss": 2.7106366005672773,
+        "val_accuracy": 0.44179999828338623
     },
     {
-        "avg_train_loss": 0.3076923076923077,
-        "train_accuracy": 0.6111111111111112,
-        "avg_val_loss": 0.3333333333333333,
-        "val_accuracy": 0.55
+        "avg_train_loss": 1.955078414747979,
+        "train_accuracy": 0.64364,
+        "avg_val_loss": 2.8602050003732087,
+        "val_accuracy": 0.421099990606308
     },
     {
-        "avg_train_loss": 0.2857142857142857,
-        "train_accuracy": 0.631578947368421,
-        "avg_val_loss": 0.3076923076923077,
-        "val_accuracy": 0.5714285714285714
+        "avg_train_loss": 1.8526526395891694,
+        "train_accuracy": 0.68184,
+        "avg_val_loss": 2.723868011668989,
+        "val_accuracy": 0.44020000100135803
     },
     {
-        "avg_train_loss": 0.26666666666666666,
-        "train_accuracy": 0.65,
-        "avg_val_loss": 0.2857142857142857,
-        "val_accuracy": 0.5909090909090909
+        "avg_train_loss": 1.7414095465830328,
+        "train_accuracy": 0.7196,
+        "avg_val_loss": 2.8222216952378583,
+        "val_accuracy": 0.4287000000476837
     },
     {
-        "avg_train_loss": 0.25,
-        "train_accuracy": 0.6666666666666666,
-        "avg_val_loss": 0.26666666666666666,
-        "val_accuracy": 0.6086956521739131
+        "avg_train_loss": 1.6265801092942251,
+        "train_accuracy": 0.7615,
+        "avg_val_loss": 2.75775924002289,
+        "val_accuracy": 0.430400013923645
     },
     {
-        "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
+        "avg_train_loss": 1.5103181164308914,
+        "train_accuracy": 0.80724,
+        "avg_val_loss": 2.8081995484175954,
+        "val_accuracy": 0.43479999899864197
     }
 ]

train.py

@@ -9,6 +9,7 @@ from PIL import Image
 from torchvision import transforms
 from torch.utils.data import DataLoader, Dataset
 from model import MyModel
+import numpy as np
 
 
 class MiniPlaces(Dataset):
@@ -75,6 +76,47 @@ class MiniPlaces(Dataset):
         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.
@@ -137,7 +179,7 @@ def train(model, train_loader, val_loader, optimizer, criterion, device,
     model = model.to(device)
 
     # Define early stopping parameters
-    patience = 3  # Number of epochs to wait for improvement
+    patience = 5  # 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
@@ -277,9 +319,12 @@ def main(args):
         transforms.Normalize(image_net_mean, image_net_std),
     ])
 
-    data_root = 'data'
+    # Separate transforms for training and validation
+    train_transform = create_train_transform()
+    val_transform = create_val_transform()
 
-    # Create MiniPlaces dataset object
+    # Create datasets
+    data_root = 'data'
     miniplaces_train = MiniPlaces(data_root,
                                   split='train',
                                   transform=data_transform)
@@ -311,6 +356,8 @@ def main(args):
     # optimizer = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-4, amsgrad=False)
     optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, dampening=0, weight_decay=1e-4, nesterov=True)
 
+    print("PARAMS NUM:", sum(p.numel() for p in model.parameters() if p.requires_grad))
+
     if args.checkpoint:
         checkpoint = torch.load(args.checkpoint)
         model.load_state_dict(checkpoint['model_state_dict'])
@@ -341,7 +388,7 @@ if __name__ == "__main__":
     parser.add_argument('--test', action='store_true')
     parser.add_argument('--checkpoint')
     parser.add_argument('--gpu', default=0)
-    parser.add_argument('--epochs', default=10)
+    parser.add_argument('--epochs', default=100)
     parser.add_argument('--batch_size', default=32)
     args = parser.parse_args()
     main(args)

train_dist.py

@@ -31,13 +31,14 @@ def setup(rank, world_size, port):
 
 def cleanup():
     """
-    Clean up the distributed training environment by destroying the process group.
+    Clean up distributed training environment
     """
-    dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()  # Synchronize all processes before destroying 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,
@@ -100,6 +101,47 @@ class MiniPlaces(Dataset):
         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.
@@ -158,146 +200,146 @@ def train_worker(rank, world_size, args):
         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)
+    try:
+        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
+
+        # 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))
+        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.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()
+                    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')
+    finally:
+        cleanup()
+        # Add explicit synchronization before exiting
+        torch.cuda.synchronize()
+        if dist.is_initialized():
+            dist.barrier()
 
 
 def test(model, test_loader, device):
@@ -345,20 +387,23 @@ def main(args):
     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)
+    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=10)
-    parser.add_argument('--batch_size', type=int, default=64)
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--batch_size', type=int, default=32)
     parser.add_argument('--port', type=int, default=4224)
     args = parser.parse_args()
     main(args)