saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-22 19:23:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add an option to export ONNX graphs in ORTModule tests (#10579)

Browse source 
Co-authored-by: Abhishek Kulkarni <abkulkarni@microsoft.com>
This commit is contained in:
Abhishek Kulkarni 2022-03-03 16:56:19 -08:00 committed by GitHub
parent 745fa5885f
commit c2c85dd6b1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 18 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py
View file

@ -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)

4
orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier_autocast.py
View file

@ -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

4
orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py
View file

@ -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)

8
orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py
View file

@ -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)

4
orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py
View file

@ -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)