Update train_mlp.py

train_mlp.py

@@ -3,9 +3,7 @@ import os
 import torch
 import torch.nn as nn
 import torch.optim as optim
-import torchvision.transforms as transforms
-from PIL import Image
-from datasets import load_dataset
+from datasets import load_from_disk
 
 # Define the MLP model
 class MLP(nn.Module):
@@ -22,22 +20,20 @@ class MLP(nn.Module):
     def forward(self, x):
         return self.model(x)
 
-# Preprocess the images
-def preprocess_image(example, image_size):
-    image = Image.open(example['image_path']).convert('RGB')
-    transform = transforms.Compose([
-        transforms.Resize((image_size, image_size)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-    image = transform(image)
-    return {'image': image, 'label': example['label']}
+# Custom collate function
+def custom_collate(batch):
+    images = torch.stack([item['image'] for item in batch])
+    labels = torch.tensor([item['label'] for item in batch])
+    return {'image': images, 'label': labels}
 
 # Train the model
-def train_model(model, train_loader, val_loader, epochs=10, lr=0.001):
+def train_model(model, train_loader, val_loader, epochs=10, lr=0.001, save_loss_path=None):
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=lr)
 
+    train_losses = []
+    val_losses = []
+
     for epoch in range(epochs):
         model.train()
         running_loss = 0.0
@@ -53,7 +49,9 @@ def train_model(model, train_loader, val_loader, epochs=10, lr=0.001):
 
             running_loss += loss.item()
 
-        print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}')
+        avg_train_loss = running_loss / len(train_loader)
+        train_losses.append(avg_train_loss)
+        print(f'Epoch {epoch+1}, Loss: {avg_train_loss}')
 
         # Validation
         model.eval()
@@ -73,9 +71,16 @@ def train_model(model, train_loader, val_loader, epochs=10, lr=0.001):
                 total += labels.size(0)
                 correct += (predicted == labels).sum().item()
 
-        print(f'Validation Loss: {val_loss/len(val_loader)}, Accuracy: {100 * correct / total}%')
+        avg_val_loss = val_loss / len(val_loader)
+        val_losses.append(avg_val_loss)
+        print(f'Validation Loss: {avg_val_loss}, Accuracy: {100 * correct / total}%')
+
+    if save_loss_path:
+        with open(save_loss_path, 'w') as f:
+            for epoch, (train_loss, val_loss) in enumerate(zip(train_losses, val_losses)):
+                f.write(f'Epoch {epoch+1}, Train Loss: {train_loss}, Validation Loss: {val_loss}\n')
 
-    return val_loss / len(val_loader)
+    return avg_val_loss
 
 # Main function
 def main():
@@ -84,21 +89,15 @@ def main():
     parser.add_argument('--width', type=int, default=512, help='Number of neurons per hidden layer (default: 512)')
     args = parser.parse_args()
 
-    # Load the dataset
-    dataset = load_dataset('your_dataset_name')
-    train_dataset = dataset['train']
-    val_dataset = dataset['validation']
+    # Load the preprocessed datasets
+    train_dataset = load_from_disk('preprocessed_train_dataset')
+    val_dataset = load_from_disk('preprocessed_val_dataset')
 
     # Determine the number of classes
     num_classes = len(set(train_dataset['label']))
 
     # Determine the fixed resolution of the images
-    example_image = Image.open(train_dataset[0]['image_path'])
-    image_size = example_image.size[0]  # Assuming the images are square
-
-    # Preprocess the dataset
-    train_dataset = train_dataset.map(lambda x: preprocess_image(x, image_size))
-    val_dataset = val_dataset.map(lambda x: preprocess_image(x, image_size))
+    image_size = train_dataset[0]['image'].size(1)  # Assuming the images are square
 
     # Define the model
     input_size = image_size * image_size * 3
@@ -107,12 +106,13 @@ def main():
 
     model = MLP(input_size, hidden_sizes, output_size)
 
-    # Create data loaders
-    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
-    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=32, shuffle=False)
+    # Create data loaders with custom collate function
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=custom_collate)
+    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=32, shuffle=False, collate_fn=custom_collate)
 
     # Train the model and get the final loss
-    final_loss = train_model(model, train_loader, val_loader)
+    save_loss_path = 'losses.txt'
+    final_loss = train_model(model, train_loader, val_loader, save_loss_path=save_loss_path)
 
     # Calculate the number of parameters
     param_count = sum(p.numel() for p in model.parameters())