Update train_mlp.py

train_mlp.py

@@ -3,7 +3,8 @@ import os
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from datasets import load_from_disk
+from datasets import load_dataset
+from torchvision import transforms
 
 # Define the MLP model
 class MLP(nn.Module):
@@ -20,12 +21,6 @@ class MLP(nn.Module):
     def forward(self, x):
         return self.model(x)
 
-# 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, save_loss_path=None):
     criterion = nn.CrossEntropyLoss()
@@ -84,20 +79,23 @@ def train_model(model, train_loader, val_loader, epochs=10, lr=0.001, save_loss_
 
 # Main function
 def main():
-    parser = argparse.ArgumentParser(description='Train an MLP on a Hugging Face dataset with JPEG images and class labels.')
+    parser = argparse.ArgumentParser(description='Train an MLP on the zh-plus/tiny-imagenet dataset.')
     parser.add_argument('--layer_count', type=int, default=2, help='Number of hidden layers (default: 2)')
     parser.add_argument('--width', type=int, default=512, help='Number of neurons per hidden layer (default: 512)')
     args = parser.parse_args()
 
-    # Load the preprocessed datasets
-    train_dataset = load_from_disk('preprocessed_train_dataset')
-    val_dataset = load_from_disk('preprocessed_val_dataset')
+    # Load the zh-plus/tiny-imagenet dataset
+    dataset = load_dataset('zh-plus/tiny-imagenet')
+
+    # Split the dataset into train and validation sets
+    train_dataset = dataset['train']
+    val_dataset = dataset['valid']
 
     # Determine the number of classes
     num_classes = len(set(train_dataset['label']))
 
     # Determine the fixed resolution of the images
-    image_size = train_dataset[0]['image'].size(1)  # Assuming the images are square
+    image_size = 64  # Assuming the images are square
 
     # Define the model
     input_size = image_size * image_size * 3
@@ -106,9 +104,18 @@ def main():
 
     model = MLP(input_size, hidden_sizes, output_size)
 
-    # 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)
+    # Define the transformation to convert PIL images to tensors
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+    ])
+
+    # Apply the transformation to the datasets
+    train_dataset.set_transform(lambda example: {'image': transform(example['image']), 'label': example['label']})
+    val_dataset.set_transform(lambda example: {'image': transform(example['image']), 'label': example['label']})
+
+    # 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)
 
     # Train the model and get the final loss
     save_loss_path = 'losses.txt'