Update train_mlp_batches.py

train_mlp_batches.py

@@ -1,3 +1,4 @@
+from modelscope.hub.api import HubApi
 import argparse
 import os
 import torch
@@ -7,9 +8,13 @@ from datasets import load_dataset
 from PIL import Image
 import numpy as np
 from torch.utils.data import DataLoader, Dataset
+from mmengine.model import BaseModel
+from mmengine.runner import Runner, EpochBasedTrainLoop, ValLoop
+from mmengine.hooks import CheckpointHook, LoggerHook
+from mmengine.optim import OptimWrapper
 
 # Define the MLP model
-class MLP(nn.Module):
+class MLP(BaseModel):
     def __init__(self, input_size, hidden_sizes, output_size):
         super(MLP, self).__init__()
         layers = []
@@ -19,9 +24,21 @@ class MLP(nn.Module):
             if i < len(sizes) - 2:
                 layers.append(nn.ReLU())
         self.model = nn.Sequential(*layers)
-
-    def forward(self, x):
-        return self.model(x)
+        self.criterion = nn.CrossEntropyLoss()
+
+    def forward(self, inputs, labels, mode='train'):
+        outputs = self.model(inputs)
+        if mode == 'train':
+            loss = self.criterion(outputs, labels)
+            return dict(loss=loss)
+        elif mode == 'val':
+            loss = self.criterion(outputs, labels)
+            _, predicted = torch.max(outputs.data, 1)
+            correct = (predicted == labels).sum().item()
+            total = labels.size(0)
+            return dict(loss=loss, correct=correct, total=total)
+        else:
+            return outputs
 
 # Custom Dataset class to handle image preprocessing
 class TinyImageNetDataset(Dataset):
@@ -40,67 +57,22 @@ class TinyImageNetDataset(Dataset):
         label = torch.tensor(example['label'])
         return img, label
 
-# Train the model
-def train_model(model, train_loader, val_loader, 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)
-
-    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
-        for batch_idx, (inputs, labels) in enumerate(train_loader):
-            inputs, labels = inputs.to(device), labels.to(device)
-
-            optimizer.zero_grad()
-            outputs = model(inputs)
-            loss = criterion(outputs, labels)
-            loss.backward()
-            optimizer.step()
-
-            running_loss += loss.item()
-
-        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()
-        val_loss = 0.0
-        correct = 0
-        total = 0
-        with torch.no_grad():
-            for inputs, labels in val_loader:
-                inputs, labels = inputs.to(device), labels.to(device)
-
-                outputs = model(inputs)
-                loss = criterion(outputs, labels)
-                val_loss += loss.item()
-
-                _, predicted = torch.max(outputs.data, 1)
-                total += labels.size(0)
-                correct += (predicted == labels).sum().item()
-
-        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}%')
-
-        # 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)):
-                f.write(f'Epoch {epoch+1}, Train Loss: {train_loss}, Validation Loss: {val_loss}\n')
-
-    return avg_val_loss
+# Define the training loop
+class MLPTrainLoop(EpochBasedTrainLoop):
+    def run_iter(self, data_batch: dict, train_mode: bool = True) -> None:
+        data_batch = self.data_preprocessor(data_batch, training=True)
+        outputs = self.model(**data_batch, mode='train')
+        parsed_outputs = self.model.parse_losses(outputs)
+        self.optim_wrapper.update_params(parsed_outputs)
+
+# Define the validation loop
+class MLPValLoop(ValLoop):
+    def run_iter(self, data_batch: dict, train_mode: bool = False) -> None:
+        data_batch = self.data_preprocessor(data_batch, training=False)
+        outputs = self.model(**data_batch, mode='val')
+        self.outputs['loss'].append(outputs['loss'].item())
+        self.outputs['correct'].append(outputs['correct'])
+        self.outputs['total'].append(outputs['total'])
 
 # Main function
 def main():
@@ -109,6 +81,7 @@ def main():
     parser.add_argument('--width', type=int, default=512, help='Number of neurons per hidden layer (default: 512)')
     parser.add_argument('--batch_size', type=int, default=8, help='Batch size for training (default: 8)')
     parser.add_argument('--save_model_dir', type=str, default='saved_models', help='Directory to save model checkpoints (default: saved_models)')
+    parser.add_argument('--access_token', type=str, required=True, help='ModelScope SDK access token')
     args = parser.parse_args()
 
     # Load the zh-plus/tiny-imagenet dataset
@@ -138,9 +111,27 @@ def main():
     train_loader = DataLoader(TinyImageNetDataset(train_dataset), batch_size=args.batch_size, shuffle=True)
     val_loader = DataLoader(TinyImageNetDataset(val_dataset), batch_size=args.batch_size, shuffle=False)
 
-    # Train the model and get the final loss
-    save_loss_path = 'losses.txt'
-    final_loss = train_model(model, train_loader, val_loader, save_loss_path=save_loss_path, save_model_dir=args.save_model_dir)
+    # Define the optimizer
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+    # Define the runner
+    runner = Runner(
+        model=model,
+        work_dir=args.save_model_dir,
+        train_dataloader=train_loader,
+        val_dataloader=val_loader,
+        optim_wrapper=dict(optimizer=optimizer),
+        train_loop=MLPTrainLoop,
+        val_loop=MLPValLoop,
+        val_interval=1,
+        default_hooks=dict(
+            checkpoint=dict(type=CheckpointHook, interval=1, save_best='auto'),
+            logger=dict(type=LoggerHook, interval=10)
+        )
+    )
+
+    # Start training
+    runner.train()
 
     # Calculate the number of parameters
     param_count = sum(p.numel() for p in model.parameters())
@@ -156,14 +147,26 @@ def main():
     # Write the results to a text file in the model folder
     result_path = os.path.join(model_folder, 'results.txt')
     with open(result_path, 'w') as f:
-        f.write(f'Layer Count: {args.layer_count}, Width: {args.width}, Parameter Count: {param_count}, Final Loss: {final_loss}\n')
+        f.write(f'Layer Count: {args.layer_count}, Width: {args.width}, Parameter Count: {param_count}\n')
 
     # Save a duplicate of the results in the 'results' folder
     results_folder = 'results'
     os.makedirs(results_folder, exist_ok=True)
     duplicate_result_path = os.path.join(results_folder, f'results_l{args.layer_count}w{args.width}.txt')
     with open(duplicate_result_path, 'w') as f:
-        f.write(f'Layer Count: {args.layer_count}, Width: {args.width}, Parameter Count: {param_count}, Final Loss: {final_loss}\n')
+        f.write(f'Layer Count: {args.layer_count}, Width: {args.width}, Parameter Count: {param_count}\n')
+
+    # Upload the model to ModelScope
+    api = HubApi()
+    api.login(args.access_token)
+    api.push_model(
+        model_id="puffy310/MLPScaling", 
+        model_dir=model_folder  # Local model directory, the directory must contain configuration.json
+    )
+
+    # Delete the local model directory
+    import shutil
+    shutil.rmtree(model_folder)
 
 if __name__ == '__main__':
     main()