Google Git
Sign in
apache / mxnet / c8922fedff48cf0c4f15bf0e5fe805be8be34512 / . / python / mxnet / util.py
blob: 919e18107cec7a668d1b24bca89f70cb5614c4a7 [file] [log] [blame]
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""general utility functions"""
import ctypes
import functools
import inspect
import threading
import tempfile
import platform
from contextlib import contextmanager
from struct import calcsize
from .base import (_LIB, check_call, c_str, py_str,
numeric_types, integer_types, long,
_MAX_VALUE_64_BIT_UNSIGNED_,
_MAX_VALUE_64_BIT_SIGNED_,
_MAX_VALUE_FLOAT32_REPRESENT_)
_np_ufunc_default_kwargs = {
'where': True,
'casting': 'same_kind',
'order': 'K',
'dtype': None,
'subok': True,
}
_set_np_shape_logged = False
_set_np_array_logged = False
_set_np_default_dtype_logged = False
def get_gpu_count():
size = ctypes.c_int()
check_call(_LIB.MXGetGPUCount(ctypes.byref(size)))
return size.value
def get_gpu_memory(gpu_dev_id):
free_mem = ctypes.c_uint64(0)
total_mem = ctypes.c_uint64(0)
check_call(_LIB.MXGetGPUMemoryInformation64(gpu_dev_id, ctypes.byref(free_mem), ctypes.byref(total_mem)))
return free_mem.value, total_mem.value
def set_np_shape(active):
"""Turns on/off NumPy shape semantics, in which `()` represents the shape of scalar tensors,
and tuples with `0` elements, for example, `(0,)`, `(1, 0, 2)`, represent the shapes
of zero-size tensors. This is turned off by default for keeping backward compatibility.
Please note that this is designed as an infrastructure for the incoming
MXNet-NumPy operators. Legacy operators registered in the modules
`mx.nd` and `mx.sym` are not guaranteed to behave like their counterparts
in NumPy within this semantics.
Parameters
----------
active : bool
Indicates whether to turn on/off NumPy shape semantics.
Returns
-------
A bool value indicating the previous state of NumPy shape semantics.
Example
-------
>>> import mxnet as mx
>>> prev_state = mx.set_np_shape(True)
>>> print(prev_state)
False
>>> print(mx.is_np_shape())
True
"""
global _set_np_shape_logged
if active:
if not _set_np_shape_logged:
import logging
logging.info('NumPy-shape semantics has been activated in your code. '
'This is required for creating and manipulating scalar and zero-size '
'tensors, which were not supported in MXNet before, as in the official '
'NumPy library. Please DO NOT manually deactivate this semantics while '
'using `mxnet.numpy` and `mxnet.numpy_extension` modules.')
_set_np_shape_logged = True
elif is_np_array():
raise ValueError('Deactivating NumPy shape semantics while NumPy array semantics is still'
' active is not allowed. Please consider calling `npx.reset_np()` to'
' deactivate both of them.')
prev = ctypes.c_int()
check_call(_LIB.MXSetIsNumpyShape(ctypes.c_int(active), ctypes.byref(prev)))
return bool(prev.value)
def is_np_shape():
"""Checks whether the NumPy shape semantics is currently turned on.
In NumPy shape semantics, `()` represents the shape of scalar tensors,
and tuples with `0` elements, for example, `(0,)`, `(1, 0, 2)`, represent
the shapes of zero-size tensors. This is turned off by default for keeping
backward compatibility.
In the NumPy shape semantics, `-1` indicates an unknown size. For example,
`(-1, 2, 2)` means that the size of the first dimension is unknown. Its size
may be inferred during shape inference.
Please note that this is designed as an infrastructure for the incoming
MXNet-NumPy operators. Legacy operators registered in the modules
`mx.nd` and `mx.sym` are not guaranteed to behave like their counterparts
in NumPy within this semantics.
Returns
-------
A bool value indicating whether the NumPy shape semantics is currently on.
Example
-------
>>> import mxnet as mx
>>> prev_state = mx.set_np_shape(True)
>>> print(prev_state)
False
>>> print(mx.is_np_shape())
True
"""
curr = ctypes.c_bool()
check_call(_LIB.MXIsNumpyShape(ctypes.byref(curr)))
return curr.value
class _NumpyShapeScope(object):
"""Scope for managing NumPy shape semantics.
In NumPy shape semantics, `()` represents the shape of scalar tensors,
and tuples with `0` elements, for example, `(0,)`, `(1, 0, 2)`, represent
the shapes of zero-size tensors.
Do not use this class directly. Use `np_shape(active)` instead.
Example::
with _NumpyShapeScope(True):
y = model(x)
backward([y])
"""
def __init__(self, is_np_shape): #pylint: disable=redefined-outer-name
self._enter_is_np_shape = is_np_shape
self._prev_is_np_shape = None
def __enter__(self):
if self._enter_is_np_shape is not None:
self._prev_is_np_shape = set_np_shape(self._enter_is_np_shape)
def __exit__(self, ptype, value, trace):
if self._enter_is_np_shape is not None and self._prev_is_np_shape != self._enter_is_np_shape:
set_np_shape(self._prev_is_np_shape)
def np_shape(active=True):
"""Returns an activated/deactivated NumPy shape scope to be used in 'with' statement
and captures code that needs the NumPy shape semantics, i.e. support of scalar and
zero-size tensors.
Please note that this is designed as an infrastructure for the incoming
MXNet-NumPy operators. Legacy operators registered in the modules
`mx.nd` and `mx.sym` are not guaranteed to behave like their counterparts
in NumPy even within this scope.
Parameters
----------
active : bool
Indicates whether to activate NumPy-shape semantics.
Returns
-------
_NumpyShapeScope
A scope object for wrapping the code w/ or w/o NumPy-shape semantics.
Example::
with mx.np_shape(active=True):
# A scalar tensor's shape is `()`, whose `ndim` is `0`.
scalar = mx.nd.ones(shape=())
assert scalar.shape == ()
# If NumPy shape semantics is enabled, 0 in a shape means that
# dimension contains zero elements.
data = mx.sym.var("data", shape=(0, 2, 3))
ret = mx.sym.sin(data)
arg_shapes, out_shapes, _ = ret.infer_shape()
assert arg_shapes[0] == (0, 2, 3)
assert out_shapes[0] == (0, 2, 3)
# -1 means unknown shape dimension size in the new NumPy shape definition
data = mx.sym.var("data", shape=(-1, 2, 3))
ret = mx.sym.sin(data)
arg_shapes, out_shapes, _ = ret.infer_shape_partial()
assert arg_shapes[0] == (-1, 2, 3)
assert out_shapes[0] == (-1, 2, 3)
# When a shape is completely unknown when NumPy shape semantics is on, it is
# represented as `None` in Python.
data = mx.sym.var("data")
ret = mx.sym.sin(data)
arg_shapes, out_shapes, _ = ret.infer_shape_partial()
assert arg_shapes[0] is None
assert out_shapes[0] is None
with mx.np_shape(active=False):
# 0 means unknown shape dimension size in the legacy shape definition.
data = mx.sym.var("data", shape=(0, 2, 3))
ret = mx.sym.sin(data)
arg_shapes, out_shapes, _ = ret.infer_shape_partial()
assert arg_shapes[0] == (0, 2, 3)
assert out_shapes[0] == (0, 2, 3)
# When a shape is completely unknown in the legacy mode (default), its ndim is
# equal to 0 and it is represented as `()` in Python.
data = mx.sym.var("data")
ret = mx.sym.sin(data)
arg_shapes, out_shapes, _ = ret.infer_shape_partial()
assert arg_shapes[0] == ()
assert out_shapes[0] == ()
"""
return _NumpyShapeScope(active)
def use_np_shape(func):
"""A decorator wrapping a function or class with activated NumPy-shape semantics.
When `func` is a function, this ensures that the execution of the function is scoped with NumPy
shape semantics, such as the support for zero-dim and zero size tensors. When
`func` is a class, it ensures that all the methods, static functions, and properties
of the class are executed with the NumPy shape semantics.
.. code-block:: python
import mxnet as mx
@mx.use_np_shape
def scalar_one():
return mx.nd.ones(())
print(scalar_one())
@np.use_np_shape
class ScalarTensor(object):
def __init__(self, val=None):
if val is None:
val = ScalarTensor.random().value
self._scalar = mx.nd.ones(()) * val
def __repr__(self):
print("Is __repr__ in np_shape semantics? {}!".format(str(np.is_np_shape())))
return str(self._scalar.asnumpy())
@staticmethod
def random():
val = mx.nd.random.uniform().asnumpy().item()
return ScalarTensor(val)
@property
def value(self):
print("Is value property in np_shape semantics? {}!".format(str(np.is_np_shape())))
return self._scalar.asnumpy().item()
print("Is global scope of np_shape activated? {}!".format(str(np.is_np_shape())))
scalar_tensor = ScalarTensor()
print(scalar_tensor)
Parameters
----------
func : a user-provided callable function or class to be scoped by the NumPy-shape semantics.
Returns
-------
Function or class
A function or class wrapped in the NumPy-shape scope.
"""
if inspect.isclass(func):
for name, method in inspect.getmembers(
func,
predicate=
lambda f: inspect.isfunction(f) or inspect.ismethod(f) or isinstance(f, property)):
if isinstance(method, property):
setattr(func, name, property(use_np_shape(method.__get__),
method.__set__,
method.__delattr__,
method.__doc__))
else:
setattr(func, name, use_np_shape(method))
return func
elif callable(func):
@functools.wraps(func)
def _with_np_shape(*args, **kwargs):
with np_shape(active=True):
return func(*args, **kwargs)
return _with_np_shape
else:
raise TypeError('use_np_shape can only decorate classes and callable objects, '
'while received a {}'.format(str(type(func))))
def _sanity_check_params(func_name, unsupported_params, param_dict):
for param_name in unsupported_params:
if param_name in param_dict:
raise NotImplementedError("function {} does not support parameter {}"
.format(func_name, param_name))
def set_module(module):
"""Decorator for overriding __module__ on a function or class.
Example usage::
@set_module('mxnet.numpy')
def example():
pass
assert example.__module__ == 'numpy'
"""
def decorator(func):
if module is not None:
func.__module__ = module
return func
return decorator
class _NumpyArrayScope(object):
"""Scope for managing NumPy array creation. This is often used
with `is_np_array=True` in initializer to enforce array creation
as type `mxnet.numpy.ndarray`, instead of `mx.nd.NDArray` in Gluon.
Do not use this class directly. Use `np_array(active)` instead.
"""
_current = threading.local()
def __init__(self, is_np_array): # pylint: disable=redefined-outer-name
self._old_scope = None
self._is_np_array = is_np_array
def __enter__(self):
if not hasattr(_NumpyArrayScope._current, "value"):
_NumpyArrayScope._current.value = _NumpyArrayScope(False)
self._old_scope = _NumpyArrayScope._current.value
_NumpyArrayScope._current.value = self
return self
def __exit__(self, ptype, value, trace):
assert self._old_scope
_NumpyArrayScope._current.value = self._old_scope
def np_array(active=True):
"""Returns an activated/deactivated NumPy-array scope to be used in 'with' statement
and captures code that needs the NumPy-array semantics.
Currently, this is used in Gluon to enforce array creation in `Block`s as type
`mxnet.numpy.ndarray`, instead of `mx.nd.NDArray`.
It is recommended to use the decorator `use_np_array` to decorate the classes
that need this semantics, instead of using this function in a `with` statement
unless you know exactly what has been scoped by this semantics.
Please note that this is designed as an infrastructure for the incoming
MXNet-NumPy operators. Legacy operators registered in the modules
`mx.nd` and `mx.sym` are not guaranteed to behave like their counterparts
in NumPy even within this scope.
Parameters
----------
active : bool
Indicates whether to activate NumPy-array semantics.
Returns
-------
_NumpyShapeScope
A scope object for wrapping the code w/ or w/o NumPy-shape semantics.
"""
return _NumpyArrayScope(active)
def is_np_array():
"""Checks whether the NumPy-array semantics is currently turned on.
This is currently used in Gluon for checking whether an array of type `mxnet.numpy.ndarray`
or `mx.nd.NDArray` should be created. For example, at the time when a parameter
is created in a `Block`, an `mxnet.numpy.ndarray` is created if this returns true; else
an `mx.nd.NDArray` is created.
Normally, users are not recommended to use this API directly unless you known exactly
what is going on under the hood.
Please note that this is designed as an infrastructure for the incoming
MXNet-NumPy operators. Legacy operators registered in the modules
`mx.nd` and `mx.sym` are not guaranteed to behave like their counterparts
in NumPy within this semantics.
Returns
-------
A bool value indicating whether the NumPy-array semantics is currently on.
"""
return _NumpyArrayScope._current.value._is_np_array if hasattr(
_NumpyArrayScope._current, "value") else False
def use_np_array(func):
"""A decorator wrapping Gluon `Block`s and all its methods, properties, and static functions
with the semantics of NumPy-array, which means that where ndarrays are created,
`mxnet.numpy.ndarray`s should be created, instead of legacy ndarrays of type `mx.nd.NDArray`.
For example, at the time when a parameter is created in a `Block`, an `mxnet.numpy.ndarray`
is created if it's decorated with this decorator.
.. code-block:: python
import mxnet as mx
from mxnet import gluon, nd, np
from mxnet.gluon import Parameter
class TestHybridBlock1(gluon.HybridBlock):
def __init__(self):
super(TestHybridBlock1, self).__init__()
self.w = Parameter('w', shape=(2, 2))
def forward(self, x):
return nd.dot(x, self.w.data())
x = mx.nd.ones((2, 2))
net1 = TestHybridBlock1()
net1.initialize()
out = net1.forward(x)
for _, v in net1.collect_params().items():
assert type(v.data()) is mx.nd.NDArray
assert type(out) is mx.nd.NDArray
@mx.util.use_np_array
class TestHybridBlock2(gluon.HybridBlock):
def __init__(self):
super(TestHybridBlock2, self).__init__()
self.w = Parameter('w', shape=(2, 2))
def forward(self, x):
return np.dot(x, self.w.data())
x = np.ones((2, 2))
net2 = TestHybridBlock2()
net2.initialize()
out = net2.forward(x)
for _, v in net2.collect_params().items():
print(type(v.data()))
assert type(v.data()) is np.ndarray
assert type(out) is np.ndarray
Parameters
----------
func : a user-provided callable function or class to be scoped by the NumPy-array semantics.
Returns
-------
Function or class
A function or class wrapped in the NumPy-array scope.
"""
if inspect.isclass(func):
for name, method in inspect.getmembers(
func,
predicate=
lambda f: inspect.isfunction(f) or inspect.ismethod(f) or isinstance(f, property)):
if isinstance(method, property):
setattr(func, name, property(use_np_array(method.__get__),
method.__set__,
method.__delattr__,
method.__doc__))
else:
setattr(func, name, use_np_array(method))
return func
elif callable(func):
@functools.wraps(func)
def _with_np_array(*args, **kwargs):
with np_array(active=True):
return func(*args, **kwargs)
return _with_np_array
else:
raise TypeError('use_np_array can only decorate classes and callable objects, '
'while received a {}'.format(str(type(func))))
def use_np(func):
"""A convenience decorator for wrapping user provided functions and classes in the scope of
both NumPy-shape and NumPy-array semantics, which means that ``(1)`` empty tuples ``()`` and
tuples with zeros, such as ``(0, 1)``, ``(1, 0, 2)``, will be treated as scalar tensors' shapes and
zero-size tensors' shapes in shape inference functions of operators, instead of as unknown
in legacy mode; (2) ndarrays of type :class:`mxnet.numpy.ndarray` should be created instead of
:class:`mx.nd.NDArray`.
.. code-block:: python
import mxnet as mx
from mxnet import gluon, nd, np
from mxnet.gluon import Parameter
class TestHybridBlock1(gluon.HybridBlock):
def __init__(self):
super(TestHybridBlock1, self).__init__()
self.w = Parameter('w', shape=(2, 2))
def forward(self, x):
return nd.dot(x, self.w.data()) + nd.ones((1,))
x = mx.nd.ones((2, 2))
net1 = TestHybridBlock1()
net1.initialize()
out = net1.forward(x)
for _, v in net1.collect_params().items():
assert type(v.data()) is mx.nd.NDArray
assert type(out) is mx.nd.NDArray
@mx.util.use_np
class TestHybridBlock2(gluon.HybridBlock):
def __init__(self):
super(TestHybridBlock2, self).__init__()
self.w = Parameter('w', shape=(2, 2))
def forward(self, x):
return np.dot(x, self.w.data()) + np.ones(())
x = np.ones((2, 2))
net2 = TestHybridBlock2()
net2.initialize()
out = net2.forward(x)
for _, v in net2.collect_params().items():
print(type(v.data()))
assert type(v.data()) is np.ndarray
assert type(out) is np.ndarray
Parameters
----------
func : a user-provided callable function or class to be scoped by the
NumPy-shape and NumPy-array semantics.
Returns
-------
Function or class
A function or class wrapped in the Numpy-shape and NumPy-array scope.
"""
return use_np_shape(use_np_array(func))
def np_ufunc_legal_option(key, value):
"""Checking if ufunc arguments are legal inputs
Parameters
----------
key : string
the key of the ufunc argument.
value : string
the value of the ufunc argument.
Returns
-------
legal : boolean
Whether or not the argument is a legal one. True when the key is one of the ufunc
arguments and value is an allowed value. False when the key is not one of the ufunc
arugments or the value is not an allowed value even when the key is a legal one.
"""
if key == 'where':
return True
elif key == 'casting':
return (value in set(['no', 'equiv', 'safe', 'same_kind', 'unsafe']))
elif key == 'order':
if isinstance(value, str):
return True
elif key == 'dtype':
import numpy as _np
return (value in set([_np.int8, _np.uint8, _np.int32, _np.int64,
_np.float16, _np.float32, _np.float64,
'int8', 'uint8', 'int32', 'int64',
'float16', 'float32', 'float64']))
elif key == 'subok':
return isinstance(value, bool)
return False
def wrap_np_unary_func(func):
"""A convenience decorator for wrapping numpy-compatible unary ufuncs to provide uniform
error handling.
Parameters
----------
func : a numpy-compatible unary function to be wrapped for better error handling.
Returns
-------
Function
A function wrapped with proper error handling.
"""
@functools.wraps(func)
def _wrap_np_unary_func(x, out=None, **kwargs):
if len(kwargs) != 0:
for key, value in kwargs.items():
# if argument is not in the set of ufunc arguments
if key not in _np_ufunc_default_kwargs:
raise TypeError("{} is an invalid keyword to function \'{}\'".format(key, func.__name__))
# if argument is one of the ufunc arguments, but not with the default value
if value != _np_ufunc_default_kwargs[key]:
# if the provided value of the argument is a legal option, raise NotImplementedError
if np_ufunc_legal_option(key, value):
raise NotImplementedError("{}={} is not implemented yet for operator {}"
.format(key, str(value), func.__name__))
# otherwise raise TypeError with not understood error message
raise TypeError("{}={} not understood for operator {}"
.format(key, value, func.__name__))
return func(x, out=out)
return _wrap_np_unary_func
def wrap_np_binary_func(func):
"""A convenience decorator for wrapping numpy-compatible binary ufuncs to provide uniform
error handling.
Parameters
----------
func : a numpy-compatible binary function to be wrapped for better error handling.
Returns
-------
Function
A function wrapped with proper error handling.
"""
@functools.wraps(func)
def _wrap_np_binary_func(x1, x2, out=None, **kwargs):
if len(kwargs) != 0:
for key, value in kwargs.items():
# if argument is not in the set of ufunc arguments
if key not in _np_ufunc_default_kwargs:
raise TypeError("{} is an invalid keyword to function \'{}\'".format(key, func.__name__))
# if argument is one of the ufunc arguments, but not with the default value
if value != _np_ufunc_default_kwargs[key]:
# if the provided value of the argument is a legal option, raise NotImplementedError
if np_ufunc_legal_option(key, value):
raise NotImplementedError("{}={} is not implemented yet".format(key, str(value)))
# otherwise raise TypeError with not understood error message
raise TypeError("{} {} not understood".format(key, value))
return func(x1, x2, out=out)
return _wrap_np_binary_func
def wrap_data_api_statical_func(func):
"""
A convenience decorator for wrapping data apis standardized statical functions to provide
context keyward backward compatibility
Parameters
----------
func : a numpy-compatible array statical function to be wrapped for context keyward change.
Returns
-------
Function
A function wrapped with context keyward changes.
"""
@functools.wraps(func)
def _wrap_api_creation_func(*args, **kwargs):
if len(kwargs) != 0:
correction = kwargs.pop('ddof', None)
if correction is not None:
kwargs['correction'] = correction
return func(*args, **kwargs)
return _wrap_api_creation_func
def wrap_data_api_linalg_func(func):
"""
A convenience decorator for wrapping data apis standardized linalg functions to provide
context keyward backward compatibility
Parameters
----------
func : a numpy-compatible array linalg function to be wrapped for context keyward change.
Returns
-------
Function
A function wrapped with context keyward changes.
"""
@functools.wraps(func)
def _wrap_linalg_func(*args, **kwargs):
if len(kwargs) != 0:
upper = kwargs.pop('UPLO', None)
rcond = kwargs.pop('rcond', None)
tol = kwargs.pop('tol', None)
if upper is not None:
if upper == 'U':
kwargs['upper'] = True
else:
kwargs['upper'] = False
if rcond is not None:
kwargs['rtol'] = rcond
if tol is not None:
kwargs['rtol'] = tol
return func(*args, **kwargs)
return _wrap_linalg_func
def wrap_sort_functions(func):
"""A convenience decorator for wrapping sort functions
Parameters
----------
func : a numpy-compatible array creation function to be wrapped for parameter keyword change.
Returns
-------
Function
A function wrapped with changed keywords.
"""
@functools.wraps(func)
def _wrap_sort_func(*args, **kwargs):
if len(kwargs) != 0:
kind = kwargs.pop('kind', None)
order = kwargs.pop('order', None)
if kind is not None:
kwargs['stable'] = kind == 'stable'
if order is not None:
raise NotImplementedError("order not supported here")
return func(*args, **kwargs)
return _wrap_sort_func
def wrap_ctx_to_device_func(func):
"""A convenience decorator for converting ctx to device keyward backward compatibility
Parameters
----------
func : a function to be wrapped for context keyward change.
Returns
-------
Function
A function wrapped with context keyward changes.
"""
@functools.wraps(func)
def _wrap_func_with_ctx(*args, **kwargs):
if len(kwargs) != 0:
device = kwargs.pop('ctx', None)
if device is not None:
kwargs['device'] = device
return func(*args, **kwargs)
return _wrap_func_with_ctx
# pylint: disable=exec-used
def numpy_fallback(func):
"""decorator for falling back to offical numpy for a specific function"""
def get_device(device, new_device):
if device is None:
return new_device
else:
if new_device is None:
new_device = device
assert device == new_device, f"inconsistent device {str(device)} and {str(new_device)}"
return device
def _as_official_np_array(object):
device = None
if hasattr(object, 'asnumpy'):
return object.asnumpy(), object.device
elif isinstance(object, (list, tuple)):
tmp = []
for arr in object:
new_arr, new_device = _as_official_np_array(arr)
device = get_device(device, new_device)
tmp.append(new_arr)
return object.__class__(tmp), device
elif isinstance(object, dict):
tmp = {}
for k, v in object.items():
new_v, new_device = _as_official_np_array(v)
device = get_device(device, new_device)
tmp[k] = new_v
return tmp, device
else:
return object, None
from .ndarray import from_numpy
from .numpy import array
from .device import current_device
def _as_mx_np_array(object, device=current_device()):
import numpy as _np
if isinstance(object, _np.ndarray):
try:
ret = from_numpy(object).as_np_ndarray()
except ValueError:
ret = array(object, dtype=object.dtype, device=device)
return (ret if ('cpu' in str(device)) else ret.to_device(device))
elif isinstance(object, (list, tuple)):
tmp = [_as_mx_np_array(arr, device) for arr in object]
return object.__class__(tmp)
elif isinstance(object, dict):
return {k:_as_mx_np_array(v, device) for k, v in object}
else:
return object
import re
func_name = func.__name__
func_doc = func.__doc__
func_source = inspect.getsource(func)
func_source = re.sub(r'np\.', 'onp.', func_source)
func_source = func_source.split('\n')[1:]
indentation = func_source[0].find('def')
if indentation == -1:
raise ValueError("should wrap a function")
stripped = []
for line in func_source:
stripped.append(line[indentation:])
stripped.insert(1, ' import numpy as onp')
func_source = '\n'.join(stripped)
local = {}
exec(func_source, None, local)
func = local[func_name]
func.__doc__ = func_doc
@functools.wraps(func)
def _fallback_to_official_np(*args, **kwargs):
# for every ndarray input, fallback
new_args, device0 = _as_official_np_array(args)
new_kwargs, device1 = _as_official_np_array(kwargs)
device = get_device(device0, device1)
ret = func(*new_args, **new_kwargs)
if ret is None:
raise ValueError("Only functions with return values are allowed to use this decorator")
ret = _as_mx_np_array(ret, device=device)
return ret
return _fallback_to_official_np
# pylint: enable=exec-used
def _set_np_array(active):
"""Turns on/off NumPy array semantics for the current thread in which `mxnet.numpy.ndarray`
is expected to be created, instead of the legacy `mx.nd.NDArray`.
Parameters
---------
active : bool
A boolean value indicating whether the NumPy-array semantics should be turned on or off.
Returns
-------
A bool value indicating the previous state of NumPy array semantics.
"""
global _set_np_array_logged
if active:
if not _set_np_array_logged:
import logging
logging.info('NumPy array semantics has been activated in your code. This allows you'
' to use operators from MXNet NumPy and NumPy Extension modules as well'
' as MXNet NumPy `ndarray`s.')
_set_np_array_logged = True
cur_state = is_np_array()
_NumpyArrayScope._current.value = _NumpyArrayScope(active)
return cur_state
def set_np(shape=True, array=True, dtype=False):
"""Setting NumPy shape and array semantics at the same time.
It is required to keep NumPy shape semantics active while activating NumPy array semantics.
Deactivating NumPy shape semantics while NumPy array semantics is still active is not allowed.
It is highly recommended to set these two flags to `True` at the same time to fully enable
NumPy-like behaviors. Please refer to the Examples section for a better understanding.
Parameters
----------
shape : bool
A boolean value indicating whether the NumPy-shape semantics should be turned on or off.
When this flag is set to `True`, zero-size and zero-dim shapes are all valid shapes in
shape inference process, instead of treated as unknown shapes in legacy mode.
array : bool
A boolean value indicating whether the NumPy-array semantics should be turned on or off.
When this flag is set to `True`, it enables Gluon code flow to use or generate `mxnet.numpy.ndarray`s
instead of `mxnet.ndarray.NDArray`. For example, a `Block` would create parameters of type
`mxnet.numpy.ndarray`.
dtype : bool
A boolean value indicating whether the NumPy-dtype semantics should be turned on or off.
When this flag is set to `True`, default dtype is float64.
When this flag is set to `False`, default dtype is float32.
Examples
--------
>>> import mxnet as mx
Creating zero-dim ndarray in legacy mode would fail at shape inference.
>>> mx.nd.ones(shape=())
mxnet.base.MXNetError: Operator _ones inferring shapes failed.
>>> mx.nd.ones(shape=(2, 0, 3))
mxnet.base.MXNetError: Operator _ones inferring shapes failed.
In legacy mode, Gluon layers would create parameters and outputs of type `mx.nd.NDArray`.
>>> from mxnet.gluon import nn
>>> dense = nn.Dense(2)
>>> dense.initialize()
>>> dense(mx.nd.ones(shape=(3, 2)))
[[0.01983214 0.07832371]
[0.01983214 0.07832371]
[0.01983214 0.07832371]]
<NDArray 3x2 @cpu(0)>
>>> [p.data() for p in dense.collect_params().values()]
[
[[0.0068339 0.01299825]
[0.0301265 0.04819721]]
<NDArray 2x2 @cpu(0)>,
[0. 0.]
<NDArray 2 @cpu(0)>]
When the `shape` flag is `True`, both shape inferences are successful.
>>> from mxnet import np, npx
>>> npx.set_np() # this is required to activate NumPy-like behaviors
>>> np.ones(shape=())
array(1.)
>>> np.ones(shape=(2, 0, 3))
array([], shape=(2, 0, 3))
When the `array` flag is `True`, Gluon layers would create parameters and outputs of type `mx.np.ndarray`.
>>> dense = nn.Dense(2)
>>> dense.initialize()
>>> dense(np.ones(shape=(3, 2)))
array([[0.01983214, 0.07832371],
[0.01983214, 0.07832371],
[0.01983214, 0.07832371]])
>>> [p.data() for p in dense.collect_params().values()]
[array([[0.0068339 , 0.01299825],
[0.0301265 , 0.04819721]]), array([0., 0.])]
>>> npx.set_np(dtype=True)
>>> np.ones(shape=()).dtype
dtype('float64')
"""
if not shape and array:
raise ValueError('NumPy Shape semantics is required in using NumPy array semantics.')
_set_np_array(array)
set_np_shape(shape)
set_np_default_dtype(dtype)
def reset_np():
"""Deactivate NumPy shape and array and deafult dtype semantics at the same time."""
set_np(shape=False, array=False, dtype=False)
_CUDA_SUCCESS = 0
def get_cuda_compute_capability(device):
"""Returns the cuda compute capability of the input `device`.
Parameters
----------
device : Device
GPU context whose corresponding cuda compute capability is to be retrieved.
Returns
-------
cuda_compute_capability : int
CUDA compute capability. For example, it returns 70 for CUDA arch equal to `sm_70`.
References
----------
https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549#file-cuda_check-py
"""
if device.device_type != 'gpu':
raise ValueError('Expecting a gpu context to get cuda compute capability, '
'while received device {}'.format(str(device)))
libnames = ('libcuda.so', 'libcuda.dylib', 'nvcuda.dll', 'cuda.dll')
for libname in libnames:
try:
cuda = ctypes.CDLL(libname)
except OSError:
continue
else:
break
else:
raise OSError("could not load any of: " + ' '.join(libnames))
# Some constants taken from cuda.h
cc_major = ctypes.c_int()
cc_minor = ctypes.c_int()
cuda_device = ctypes.c_int()
error_str = ctypes.c_char_p()
ret = cuda.cuInit(0)
if ret != _CUDA_SUCCESS:
cuda.cuGetErrorString(ret, ctypes.byref(error_str))
raise RuntimeError('cuInit failed with erro code {}: {}'
.format(ret, error_str.value.decode()))
ret = cuda.cuDeviceGet(ctypes.byref(cuda_device), device.device_id)
if ret != _CUDA_SUCCESS:
cuda.cuGetErrorString(ret, ctypes.byref(error_str))
raise RuntimeError('cuDeviceGet failed with error code {}: {}'
.format(ret, error_str.value.decode()))
ret = cuda.cuDeviceComputeCapability(ctypes.byref(cc_major), ctypes.byref(cc_minor), cuda_device)
if ret != _CUDA_SUCCESS:
cuda.cuGetErrorString(ret, ctypes.byref(error_str))
raise RuntimeError('cuDeviceComputeCapability failed with error code {}: {}'
.format(ret, error_str.value.decode()))
return cc_major.value * 10 + cc_minor.value
def default_array(source_array, device=None, dtype=None):
"""Creates an array from any object exposing the default(nd or np) array interface.
Parameters
----------
source_array : array_like
An object exposing the array interface, an object whose `__array__`
method returns an array, or any (nested) sequence.
device : Device, optional
Device context (default is the current default context).
dtype : str or numpy.dtype, optional
The data type of the output array. The default dtype is ``source_array.dtype``
if `source_array` is an `NDArray`, `float32` otherwise.
Returns
-------
NDArray
An `NDArray`(nd or np) with the same contents as the `source_array`.
"""
from . import nd as _mx_nd
from . import np as _mx_np
if is_np_array():
return _mx_np.array(source_array, device=device, dtype=dtype)
else:
return _mx_nd.array(source_array, ctx=device, dtype=dtype)
class _NumpyDefaultDtypeScope(object):
"""Scope for managing NumPy default dtype semantics.
In NumPy default dtype semantics, default dtype is 'float64',
i.e. np.array([1, 2, 3]).dtype = np.float64
Original default dtype without this semantic is 'float32'.
Do not use this class directly. Use `np_shape(active)` instead.
Example::
with _NumpyDefaultDtypeScope(True):
y = model(x)
backward([y])
"""
def __init__(self, is_np_default_dtype): #pylint: disable=redefined-outer-name
self._enter_is_np_default_dtype = is_np_default_dtype
self._prev_is_np_default_dtype = None
def __enter__(self):
if self._enter_is_np_default_dtype is not None:
self._prev_is_np_default_dtype = set_np_default_dtype(self._enter_is_np_default_dtype)
def __exit__(self, ptype, value, trace):
if self._enter_is_np_default_dtype is not None and\
self._prev_is_np_default_dtype != self._enter_is_np_default_dtype:
set_np_default_dtype(self._prev_is_np_default_dtype)
def np_default_dtype(active=True):
"""Returns an activated/deactivated NumPy-default_dtype scope to be used in 'with' statement
and captures code that needs the NumPy default dtype semantics. i.e. default dtype is float64.
Please note that this is designed as an infrastructure for the incoming
MXNet-NumPy operators. Legacy operators registered in the modules
`mx.nd` and `mx.sym` are not guaranteed to behave like their counterparts
in NumPy even within this scope.
Parameters
----------
active : bool
Indicates whether to activate NumPy default dtype semantics.
Returns
-------
_NumpyDefaultDtypeScope
A scope object for wrapping the code w/ or w/o NumPy-default_dtype semantics.
Example::
with mx.np_default_dtype(active=True):
# Default Dtype is 'float64', consistent with offical NumPy behavior.
arr = mx.np.array([1, 2, 3])
assert arr.dtype == 'float64'
with mx.np_default_dtype(active=False):
# Default Dtype is 'float32' in the legacy default dtype definition.
arr = mx.np.array([1, 2, 3])
assert arr.dtype == 'float32'
"""
return _NumpyDefaultDtypeScope(active)
def use_np_default_dtype(func):
"""A decorator wrapping a function or class with activated NumPy-default_dtype semantics.
When `func` is a function, this ensures that the execution of the function is scoped with NumPy
default dtype semantics, with the support for float64 as default dtype.
When`func` is a class, it ensures that all the methods, static functions, and properties
of the class are executed with the NumPy-default_dtype semantics.
.. code-block:: python
import mxnet as mx
@mx.use_np_default_dtype
def float64_one():
return mx.nd.ones(()).dtype
print(float64_one())
@np.use_np_default_dtype
class Float64Tensor(object):
def __init__(self, data=None):
if data is None:
data = Float64Tensor.random().data
self._data = data
def __repr__(self):
print("Is __repr__ in np_default_dtype semantics? {}!".format(str(np.is_np_deafult_dtype())))
return str(self._data.asnumpy())
@staticmethod
def random():
data = mx.nd.random.uniform(shape=(2,2))
return ScalarTensor(data)
@property
def value(self):
print("Is value property in np_dafault_dtype semantics? {}!".format(str(np.is_np_default_dtype())))
return self._data.asnumpy()
print("Is global scope of np_default_dtype activated? {}!".format(str(np.is_np_default_dtype())))
float64_tensor = Float64Tensor()
print(float64_tensor)
Parameters
----------
func : a user-provided callable function or class to be scoped by the NumPy-default_dtype semantics.
Returns
-------
Function or class
A function or class wrapped in the NumPy-default_dtype scope.
"""
if inspect.isclass(func):
for name, method in inspect.getmembers(
func,
predicate=
lambda f: inspect.isfunction(f) or inspect.ismethod(f) or isinstance(f, property)):
if isinstance(method, property):
setattr(func, name, property(use_np_default_dtype(method.__get__),
method.__set__,
method.__delattr__,
method.__doc__))
else:
setattr(func, name, use_np_default_dtype(method))
return func
elif callable(func):
@functools.wraps(func)
def _with_np_default_dtype(*args, **kwargs):
with np_default_dtype(active=True):
return func(*args, **kwargs)
return _with_np_default_dtype
else:
raise TypeError('use_np_default_dtype can only decorate classes and callable objects, '
'while received a {}'.format(str(type(func))))
def is_np_default_dtype():
"""Checks whether the NumPy default dtype semantics is currently turned on.
In NumPy default dtype semantics, default dtype is float64.
Please note that this is designed as an infrastructure for the incoming
MXNet-NumPy operators. Legacy operators registered in the modules
`mx.nd` and `mx.sym` are not guaranteed to behave like their counterparts
in NumPy even within this scope.
Returns
-------
A bool value indicating whether the NumPy default dtype semantics is currently on.
See Also
--------
set_np_default_dtype : Set default dtype equals to offical numpy
set_np : npx.set_np(dtype=True) has equal performance to npx.set_np_default_dtype(True)
Example
-------
>>> import mxnet as mx
>>> from mxnet import npx
>>> prev_state = npx.set_np_default_dtype(True)
>>> print(prev_state)
False
>>> print(npx.is_np_default_dtype())
True
"""
curr = ctypes.c_bool()
check_call(_LIB.MXIsNumpyDefaultDtype(ctypes.byref(curr)))
return curr.value
def set_np_default_dtype(is_np_default_dtype=True): # pylint: disable=redefined-outer-name
"""Turns on/off NumPy default dtype semantics, because mxnet.numpy.ndarray use
32 bit data storage as default (e.g. float32 and int 32) while offical NumPy use
64 bit data storage as default (e.g. float64 and int64).
This is turned off by default for keeping backward compatibility.
Please note that this is designed as an infrastructure for the incoming
MXNet-NumPy operators. Legacy operators registered in the modules
`mx.nd` and `mx.sym` are not guaranteed to behave like their counterparts
in NumPy within this semantics.
Parameters
----------
active : bool
Indicates whether to turn on/off NumPy default dtype semantics.
Returns
-------
A bool value indicating the previous state of NumPy default dtype semantics.
Example
-------
>>> import mxnet as mx
>>> from mxnet import npx
>>> prev_state = npx.set_np_default_dtype(True)
>>> print(prev_state)
False
>>> print(npx.is_np_default_dtype())
True
"""
global _set_np_default_dtype_logged
if is_np_default_dtype:
if not _set_np_default_dtype_logged:
import logging
logging.info('NumPy array default dtype has been changed from flaot32 to float64 in your code.')
_set_np_default_dtype_logged = True
prev = ctypes.c_bool()
check_call(_LIB.MXSetIsNumpyDefaultDtype(ctypes.c_bool(is_np_default_dtype), ctypes.byref(prev)))
return prev.value
def getenv(name):
"""Get the setting of an environment variable from the C Runtime.
Parameters
----------
name : string type
The environment variable name
Returns
-------
value : string
The value of the environment variable, or None if not set
"""
ret = ctypes.c_char_p()
check_call(_LIB.MXGetEnv(c_str(name), ctypes.byref(ret)))
return None if ret.value is None else py_str(ret.value)
def setenv(name, value):
"""Set an environment variable in the C Runtime.
Parameters
----------
name : string type
The environment variable name
value : string type
The desired value to set the environment value to
"""
passed_value = None if value is None else c_str(value)
check_call(_LIB.MXSetEnv(c_str(name), passed_value))
def get_max_supported_compute_capability():
"""Get the maximum compute capability (SM arch) supported by the nvrtc compiler
"""
max_supported_cc = ctypes.c_int()
check_call(_LIB.MXGetMaxSupportedArch(ctypes.byref(max_supported_cc)))
return max_supported_cc.value
def get_rtc_compile_opts(device):
"""Get the compile ops suitable for the context, given the toolkit/driver config
"""
device_cc = get_cuda_compute_capability(device)
max_supported_cc = get_max_supported_compute_capability()
# CUDA toolkits starting with 11.1 (first to support arch 86) can compile directly to SASS
can_compile_to_SASS = max_supported_cc >= 86
should_compile_to_SASS = can_compile_to_SASS and \
device_cc <= max_supported_cc
device_cc_as_used = min(device_cc, max_supported_cc)
arch_opt = "--gpu-architecture={}_{}".format("sm" if should_compile_to_SASS else "compute",
device_cc_as_used)
return [arch_opt]
def set_flush_denorms(value):
"""Change floating-point calculations on CPU when dealing with denormalized values.
This is only applicable to architectures which supports flush-to-zero.
Denormalized values are positive and negative values that are very close to 0
(exponent is the smallest possible value).
Flushing denormalized values to 0 can speedup calculations if such values occurs,
but if fulfilling whole IEEE 754 standard is required this option should be disabled.
Flushing denormalized values is enabled in MXNet by default.
Parameters
----------
value : bool
State of flush-to-zero and denormals-are-zero in MXCSR register
Returns
-------
prev_state : bool
Previous state of flush-to-zero in MXCSR register
"""
ret = ctypes.c_bool()
passed_value = ctypes.c_bool(value)
check_call(_LIB.MXSetFlushDenorms(passed_value, ctypes.byref(ret)))
return ret.value
def dtype_from_number(number):
"""Get the data type from the given int or float number
"""
assert isinstance(number, numeric_types),\
"The input number should be either int for float types"
import numpy as _np
if isinstance(number, (int, long)):
if number > _MAX_VALUE_64_BIT_UNSIGNED_:
raise OverflowError("Integer out of bounds")
if number > _MAX_VALUE_64_BIT_SIGNED_:
return _np.uint64
elif calcsize("P") == 8:
return _np.int64
else:
return _np.int32
elif isinstance(number, float):
if abs(number) > _MAX_VALUE_FLOAT32_REPRESENT_ or \
((not _np.isnan(number)) and \
(_np.float32(number) == int(number)) and \
(number != int(number))):
return _np.float64
else:
return _np.float64 if is_np_default_dtype() else _np.float32
elif isinstance(number, _np.generic):
return number.dtype
raise TypeError('type {} not supported'.format(str(type(number))))
# This is a wrapping of tempfile.TemporaryDirectory(), known to have cleanup issues on Windows.
# The problem is partially handled as of Python 3.10 by the adding of a 'ignore_cleanup_errors'
# parameter. Once MXNet's Python version is forced to be >= 3.10, a simplification of this
# function to use 'ignore_cleanup_errors' would be possible. Until the fundamental Windows
# issues are resolved, best to use this routine instead of tempfile.TemporaryDirectory().
@contextmanager
def TemporaryDirectory(*args, **kwargs):
"""A context wrapper of tempfile.TemporaryDirectory() that ignores cleanup errors on Windows.
"""
dir = tempfile.TemporaryDirectory(*args, **kwargs)
try:
yield dir.name
finally:
try:
dir.cleanup()
except PermissionError:
if platform.system() != 'Windows':
raise