diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py index 5d7508d3b8..7524c0f382 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py @@ -323,6 +323,8 @@ def main(): help='input batch size for testing (default: 64)') parser.add_argument('--view-graphs', action='store_true', default=False, help='views forward and backward graphs') + parser.add_argument('--export-onnx-graphs', action='store_true', default=False, + help='export ONNX graphs to current directory') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument('--epochs', type=int, default=4, metavar='N', @@ -377,7 +379,7 @@ def main(): if not args.pytorch_only: # Just for future debugging - debug_options = DebugOptions(save_onnx=False, onnx_prefix='BertForSequenceClassification') + debug_options = DebugOptions(save_onnx=args.export_onnx_graphs, onnx_prefix='BertForSequenceClassification') model = ORTModule(model, debug_options) diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier_autocast.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier_autocast.py index ddc524a665..dcbe45ff19 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier_autocast.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier_autocast.py @@ -328,6 +328,8 @@ def main(): help='input batch size for testing (default: 64)') parser.add_argument('--view-graphs', action='store_true', default=False, help='views forward and backward graphs') + parser.add_argument('--export-onnx-graphs', action='store_true', default=False, + help='export ONNX graphs to current directory') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument('--epochs', type=int, default=4, metavar='N', @@ -390,7 +392,7 @@ def main(): if not args.pytorch_only: # Just for future debugging - debug_options = DebugOptions(save_onnx=False, onnx_prefix='BertForSequenceClassificationAutoCast') + debug_options = DebugOptions(save_onnx=args.export_onnx_graphs, onnx_prefix='BertForSequenceClassificationAutoCast') model = ORTModule(model, debug_options) model._torch_module._execution_manager(is_training=True)._enable_grad_acc_optimization = True diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py index 2b12b7d67c..fde4397216 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py @@ -140,6 +140,8 @@ def main(): help='how many batches to wait before logging training status (default: 300)') parser.add_argument('--view-graphs', action='store_true', default=False, help='views forward and backward graphs') + parser.add_argument('--export-onnx-graphs', action='store_true', default=False, + help='export ONNX graphs to current directory') parser.add_argument('--epochs', type=int, default=10, metavar='N', help='number of epochs to train (default: 10)') parser.add_argument('--log-level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], default='WARNING', @@ -199,7 +201,7 @@ def main(): log_level = log_level_mapping.get(args.log_level.upper(), None) if not isinstance(log_level, LogLevel): raise ValueError('Invalid log level: %s' % args.log_level) - debug_options = DebugOptions(log_level=log_level, save_onnx=False, onnx_prefix='MNIST') + debug_options = DebugOptions(log_level=log_level, save_onnx=args.export_onnx_graphs, onnx_prefix='MNIST') model = ORTModule(model, debug_options) diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py index 2ba77a6a5a..f8c33742a0 100755 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py @@ -9,7 +9,7 @@ from torchvision import datasets, transforms import time from torch.nn.parallel import DistributedDataParallel as DDP import os -from onnxruntime.training.ortmodule import ORTModule +from onnxruntime.training.ortmodule import ORTModule, DebugOptions import numpy as np # Usage : @@ -179,7 +179,9 @@ def train( if args.use_ortmodule: print("Converting to ORTModule....") - model = ORTModule(model) + debug_options = DebugOptions(save_onnx=args.export_onnx_graphs, onnx_prefix='NeuralNet') + + model = ORTModule(model, debug_options) train_dataloader, test_dataloader = get_dataloader(args,rank,args.batch_size) loss_fn = my_loss @@ -239,6 +241,8 @@ if __name__ == "__main__": help='number of steps to train. Set -1 to run through whole dataset (default: -1)') parser.add_argument('--view-graphs', action='store_true', default=False, help='views forward and backward graphs') + parser.add_argument('--export-onnx-graphs', action='store_true', default=False, + help='export ONNX graphs to current directory') parser.add_argument('--log-interval', type=int, default=300, metavar='N', help='how many batches to wait before logging training status (default: 300)') parser.add_argument("--cpu", action="store_true", default=False) diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py index 8aeab01308..4727dcbb44 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py @@ -131,6 +131,8 @@ def main(): help='how many batches to wait before logging training status (default: 300)') parser.add_argument('--view-graphs', action='store_true', default=False, help='views forward and backward graphs') + parser.add_argument('--export-onnx-graphs', action='store_true', default=False, + help='export ONNX graphs to current directory') parser.add_argument('--epochs', type=int, default=5, metavar='N', help='number of epochs to train (default: 10)') parser.add_argument('--log-level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], default='WARNING', @@ -169,7 +171,7 @@ def main(): print('Training MNIST on ORTModule....') # Just for future debugging - debug_options = DebugOptions(save_onnx=False, onnx_prefix='MNIST') + debug_options = DebugOptions(save_onnx=args.export_onnx_graphs, onnx_prefix='MNIST') model = ORTModule(model, debug_options)