Update train_mlp.py

train_mlp.py

@@ -23,7 +23,7 @@ class MLP(nn.Module):
         return self.model(x)
 
 # Train the model
-def train_model(model, train_dataset, val_dataset, epochs=10, lr=0.001, save_loss_path=None):
+def train_model(model, train_dataset, val_dataset, epochs=10, lr=0.001, save_loss_path=None, save_model_dir=None):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
 
@@ -80,6 +80,11 @@ def train_model(model, train_dataset, val_dataset, epochs=10, lr=0.001, save_los
         val_losses.append(avg_val_loss)
         print(f'Validation Loss: {avg_val_loss}, Accuracy: {100 * correct / total}%')
 
+        # Save the model after each epoch
+        if save_model_dir:
+            model_path = os.path.join(save_model_dir, f'model_epoch_{epoch+1}.pth')
+            torch.save(model.state_dict(), model_path)
+
     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)):
@@ -92,6 +97,7 @@ def main():
     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)')
+    parser.add_argument('--save_model_dir', type=str, default='saved_models', help='Directory to save model checkpoints (default: saved_models)')
     args = parser.parse_args()
 
     # Load the zh-plus/tiny-imagenet dataset
@@ -114,9 +120,12 @@ def main():
 
     model = MLP(input_size, hidden_sizes, output_size)
 
+    # Create the directory to save models
+    os.makedirs(args.save_model_dir, exist_ok=True)
+
     # Train the model and get the final loss
     save_loss_path = 'losses.txt'
-    final_loss = train_model(model, train_dataset, val_dataset, save_loss_path=save_loss_path)
+    final_loss = train_model(model, train_dataset, val_dataset, save_loss_path=save_loss_path, save_model_dir=args.save_model_dir)
 
     # Calculate the number of parameters
     param_count = sum(p.numel() for p in model.parameters())
@@ -125,7 +134,7 @@ def main():
     model_folder = f'mlp_model_l{args.layer_count}w{args.width}'
     os.makedirs(model_folder, exist_ok=True)
 
-    # Save the model
+    # Save the final model
     model_path = os.path.join(model_folder, 'model.pth')
     torch.save(model.state_dict(), model_path)