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

Suppress tracer warnings from onnx export in ORTModule (#7221)

Browse source 
* Suppress tracer warnings from onnx export in ORTModule
This commit is contained in:
baijumeswani 2021-04-08 03:41:38 -07:00 committed by GitHub
parent 27e778909d
commit d272c8434d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 59 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
orttraining/orttraining/python/training/_ortmodule_graph_execution_manager.py
View file

@ -4,6 +4,7 @@
# --------------------------------------------------------------------------
from . import _utils, _ortmodule_utils, _ortmodule_output_transformation as _ortmodule_io
from . import _ortmodule_logger as _logger
from onnxruntime.capi.onnxruntime_inference_collection import OrtValue
from onnxruntime.capi import _pybind_state as C
@ -92,8 +93,8 @@ class GraphExecutionManager(ABC):
self._input_names_require_grad = None
self._module_output_schema = None
# Verbosity for logging
self._verbosity = _ortmodule_utils.Verbosity.WARNING
# Log level
self._loglevel = _logger.LogLevel.WARNING
# TODO: Single device support for now
self._device = _utils.get_device_from_module(module)
@ -113,7 +114,7 @@ class GraphExecutionManager(ABC):
self._use_external_gpu_allocator = True
if self._use_external_gpu_allocator:
# CPP extension to get torch GPU allocator's alloc and free function addresses
self._torch_gpu_allocator = _ortmodule_utils._load_torch_gpu_allocator_cpp_extension(self._verbosity,
self._torch_gpu_allocator = _ortmodule_utils._load_torch_gpu_allocator_cpp_extension(self._loglevel < _logger.LogLevel.WARNING,
self.is_rocm_pytorch)
self._torch_alloc = self._torch_gpu_allocator.gpu_caching_allocator_raw_alloc_address()
self._torch_free = self._torch_gpu_allocator.gpu_caching_allocator_raw_delete_address()
@ -164,7 +165,7 @@ class GraphExecutionManager(ABC):
# default to PRIORITY_BASED execution order
session_options.execution_order = onnxruntime.ExecutionOrder.PRIORITY_BASED
# 0:Verbose, 1:Info, 2:Warning. 3:Error, 4:Fatal. Default is 2.
session_options.log_severity_level = int(self._verbosity)
session_options.log_severity_level = int(self._loglevel)
# enable dumping optimized training graph
if self._save_onnx:
@ -216,7 +217,7 @@ class GraphExecutionManager(ABC):
assert self._export_mode is not None, "Please use a concrete instance of ExecutionManager"
try:
with torch.no_grad():
with torch.no_grad(), _logger.suppress_os_stream_output(log_level=self._loglevel):
torch.onnx.export(self._flattened_module,
sample_inputs_copy + (sample_kwargs_copy, ),
f,
@ -226,7 +227,7 @@ class GraphExecutionManager(ABC):
do_constant_folding=False,
training=self._export_mode,
dynamic_axes=dynamic_axes,
verbose=self._verbosity < _ortmodule_utils.Verbosity.WARNING,
verbose=self._loglevel < _logger.LogLevel.WARNING,
export_params=False,
keep_initializers_as_inputs=True)
except RuntimeError as e:

51
orttraining/orttraining/python/training/_ortmodule_logger.py Normal file
View file

@ -0,0 +1,51 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from contextlib import contextmanager
from enum import IntEnum
import io
import sys
import warnings
class LogLevel(IntEnum):
VERBOSE = 0
INFO = 1
WARNING = 2
ERROR = 3
FATAL = 4
@contextmanager
def suppress_os_stream_output(suppress_stdout=True, suppress_stderr=True, log_level=LogLevel.WARNING):
"""Supress output from being printed to stdout and stderr if log_level is WARNING or higher.
If there is any output detected, a single warning is issued at of the context
"""
# stdout and stderr is written to a tempfile instead
stdout = sys.stdout
stderr = sys.stderr
suppress_logs = log_level >= LogLevel.WARNING
fo = io.StringIO()
try:
if suppress_stdout and suppress_logs:
sys.stdout = fo
if suppress_stderr and suppress_logs:
sys.stderr = fo
yield
finally:
if suppress_stdout:
sys.stdout = stdout
if suppress_stderr:
sys.stderr = stderr
if fo.tell() > 0 and suppress_logs:
# If anything was captured in fo, raise a single user warning letting users know that there was
# some warning or error that was raised
warnings.warn("There were one or more warnings or errors raised while exporting the PyTorch "
"model. Please enable INFO level logging to view all warnings and errors.", UserWarning)

12
orttraining/orttraining/python/training/_ortmodule_utils.py
View file

@ -13,15 +13,6 @@ import torch
from torch.utils.dlpack import from_dlpack, to_dlpack
from torch.utils.cpp_extension import load_inline
from enum import IntEnum
class Verbosity(IntEnum):
VERBOSE = 0
INFO = 1
WARNING = 2
ERROR = 3
FATAL = 4
def _ortvalue_to_torch_tensor(ortvalue):
# PyTorch's to_dlpack() uses same config for both torch.bool and torch.uint8,
# and convert the config to torch.uint8 tensor duing from_dlpack().
@ -29,7 +20,6 @@ def _ortvalue_to_torch_tensor(ortvalue):
torch_tensor = from_dlpack(ortvalue._ortvalue.to_dlpack())
return torch_tensor.to(torch.bool) if ortvalue.data_type() == 'tensor(bool)' else torch_tensor
def _ortvalue_from_torch_tensor(torch_tensor):
return OrtValue(C.OrtValue.from_dlpack(to_dlpack(torch_tensor), torch_tensor.dtype == torch.bool))
@ -49,7 +39,7 @@ def _load_torch_gpu_allocator_cpp_extension(verbosity, is_rocm_pytorch):
extra_cflags=['-D__HIP_PLATFORM_HCC__=1' if is_rocm_pytorch else ''],
functions=['gpu_caching_allocator_raw_alloc_address',
'gpu_caching_allocator_raw_delete_address'],
verbose=verbosity < Verbosity.WARNING, with_cuda=True)
verbose=verbosity, with_cuda=True)
def _check_same_device(device, argument_str, *args):
'''Check that all tensor arguments in *args reside on the same device as the input device'''