adding augmentation and different datasets for test and train

train.py

@@ -6,26 +6,40 @@ import torchvision.transforms as transforms
 from torch.utils.data import DataLoader, Subset
 from model import get_model, save_model
 from tqdm import tqdm
+import os
+from datetime import datetime
 
 def get_transforms():
     """
-    Define the image transformations
+    Define the image transformations with augmentation for training
     """
-    return transforms.Compose([
+    train_transform = transforms.Compose([
         transforms.Resize(224),
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomRotation(15),
+        transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
+        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                            std=[0.229, 0.224, 0.225])
     ])
+    
+    test_transform = transforms.Compose([
+        transforms.Resize(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
+                           std=[0.229, 0.224, 0.225])
+    ])
+    
+    return train_transform, test_transform
 
 def get_data(subset_size=None, train=True):
     """
     Load and prepare the dataset
-    Args:
-        subset_size (int): If provided, return only a subset of data
-        train (bool): If True, return training data, else test data
     """
-    transform = get_transforms()
+    train_transform, test_transform = get_transforms()
+    transform = train_transform if train else test_transform
+    
     dataset = torchvision.datasets.CIFAR10(
         root='./data', 
         train=train,
@@ -66,26 +80,34 @@ def evaluate_model(model, testloader, device):
 
 def train_model(model, trainloader, testloader, epochs=100, device='cuda'):
     """
-    Train the model
-    Args:
-        model: The ResNet50 model
-        trainloader: DataLoader for training data
-        testloader: DataLoader for test data
-        epochs (int): Number of epochs to train
-        device (str): Device to train on ('cuda' or 'cpu')
+    Train the model with improved hyperparameters and markdown logging
     """
     model = model.to(device)
     criterion = nn.CrossEntropyLoss()
-    optimizer = optim.Adam(model.parameters(), lr=0.001)
-    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
-        optimizer, 
-        'max', 
-        patience=5
+    
+    # Add weight decay and reduce initial learning rate
+    optimizer = optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.01)
+    
+    # Modify scheduler for better learning rate adjustment
+    scheduler = optim.lr_scheduler.OneCycleLR(
+        optimizer,
+        max_lr=0.001,
+        epochs=epochs,
+        steps_per_epoch=len(trainloader),
+        pct_start=0.2  # Warm up for first 20% of training
     )
     
-    best_acc = 0.0
+    # Create a markdown file for logging
+    log_dir = 'logs'
+    os.makedirs(log_dir, exist_ok=True)
+    log_file = os.path.join(log_dir, f'training_log_{datetime.now().strftime("%Y%m%d_%H%M%S")}.md')
     
-    # Create epoch progress bar without a description (we'll use it for stats only)
+    with open(log_file, 'w') as f:
+        f.write("# Training Log\n\n")
+        f.write("| Epoch | Train Loss | Train Acc | Test Acc | Best Acc |\n")
+        f.write("|-------|------------|-----------|-----------|----------|\n")
+
+    best_acc = 0.0
     epoch_pbar = tqdm(range(epochs), desc='Training Progress', position=0)
     
     for epoch in epoch_pbar:
@@ -108,6 +130,7 @@ def train_model(model, trainloader, testloader, epochs=100, device='cuda'):
             loss = criterion(outputs, labels)
             loss.backward()
             optimizer.step()
+            scheduler.step()  # Step the scheduler every batch
             
             running_loss += loss.item()
             _, predicted = outputs.max(1)
@@ -124,25 +147,39 @@ def train_model(model, trainloader, testloader, epochs=100, device='cuda'):
         test_acc = evaluate_model(model, testloader, device)
         epoch_pbar.write(f'Epoch {epoch+1}: Train Loss: {avg_loss:.3f} | Train Acc: {epoch_acc:.2f}% | Test Acc: {test_acc:.2f}%')
         
-        scheduler.step(test_acc)  # Using test accuracy for scheduler
-        
+        # After computing metrics, log to markdown file
+        with open(log_file, 'a') as f:
+            f.write(f"| {epoch+1:5d} | {avg_loss:.3f} | {epoch_acc:.2f}% | {test_acc:.2f}% | {best_acc:.2f}% |\n")
+
         if test_acc > best_acc:
             best_acc = test_acc
             save_model(model, 'best_model.pth')
             epoch_pbar.write(f'New best test accuracy: {test_acc:.2f}%')
+            # Add a marker for best accuracy in the markdown
+            with open(log_file, 'a') as f:
+                f.write(f"**New best accuracy achieved at epoch {epoch+1}**\n\n")
         
         if test_acc > 70:
             epoch_pbar.write(f"\nReached target accuracy of 70% on test data!")
+            with open(log_file, 'a') as f:
+                f.write(f"\n**Training stopped at epoch {epoch+1} after reaching target accuracy of 70%**\n")
             break
 
+    # Add final summary to markdown
+    with open(log_file, 'a') as f:
+        f.write(f"\n## Training Summary\n")
+        f.write(f"- Final Test Accuracy: {test_acc:.2f}%\n")
+        f.write(f"- Best Test Accuracy: {best_acc:.2f}%\n")
+        f.write(f"- Total Epochs: {epoch+1}\n")
+
 if __name__ == "__main__":
     # Set device
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
     
-    # Get train and test data
-    trainloader = get_data(subset_size=5000, train=True)
-    testloader = get_data(subset_size=1000, train=False)
+    # Get train and test data with larger batch size
+    trainloader = get_data(subset_size=10000, train=True)  # Increased from 5000
+    testloader = get_data(subset_size=2000, train=False)   # Increased from 1000
     
     # Initialize model
     model = get_model(num_classes=10)