Update train_mlp_batches.py

train_mlp_batches.py

@@ -82,8 +82,13 @@ def main():
     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')
+    parser.add_argument('--upload_checkpoint', action='store_true', help='Upload checkpoint to ModelScope')
+    parser.add_argument('--delete_checkpoint', action='store_true', help='Delete local checkpoint after uploading')
     args = parser.parse_args()
 
+    # Set up Git to use hf-mirror as a proxy
+    os.environ['GIT_PROXY_COMMAND'] = 'proxychains4 git'
+
     # Load the zh-plus/tiny-imagenet dataset
     dataset = load_dataset('zh-plus/tiny-imagenet')
 
@@ -125,7 +130,7 @@ def main():
         val_loop=MLPValLoop,
         val_interval=1,
         default_hooks=dict(
-            checkpoint=dict(type=CheckpointHook, interval=1, save_best='auto'),
+            checkpoint=dict(type=CheckpointHook, interval=1, save_best='auto') if not args.delete_checkpoint else None,
             logger=dict(type=LoggerHook, interval=10)
         )
     )
@@ -156,17 +161,19 @@ def main():
     with open(duplicate_result_path, 'w') as f:
         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
-    )
+    # Upload the model to ModelScope if specified
+    if args.upload_checkpoint:
+        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)
+    # Delete the local model directory if specified
+    if args.delete_checkpoint:
+        import shutil
+        shutil.rmtree(model_folder)
 
 if __name__ == '__main__':
     main()