Source code for mxnet.util
# 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 os
import sys
import functools
import itertools
import inspect
import threading
from .base import _LIB, check_call
_np_ufunc_default_kwargs = {
'where': True,
'casting': 'same_kind',
'order': 'K',
'dtype': None,
'subok': True,
}
def makedirs(d):
"""Create directories recursively if they don't exist. os.makedirs(exist_ok=True) is not
available in Python2"""
if sys.version_info[0] < 3:
from distutils.dir_util import mkpath
mkpath(d)
else:
os.makedirs(d, exist_ok=True) # pylint: disable=unexpected-keyword-arg
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
"""
if active:
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.')
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 wraps_safely(wrapped, assigned=functools.WRAPPER_ASSIGNMENTS):
"""This function is safe version of `functools.wraps` in Python2 which skips wrapping functions
for the attributes that do not exist."""
if sys.version_info[0] > 2:
return functools.wraps(wrapped)
else:
return functools.wraps(wrapped,
assigned=itertools.ifilter(
functools.partial(hasattr, wrapped), assigned))
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.
Example::
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):
@wraps_safely(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)
[docs]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.
Example::
import mxnet as mx
from mxnet import gluon, np
class TestHybridBlock1(gluon.HybridBlock):
def __init__(self):
super(TestHybridBlock1, self).__init__()
self.w = self.params.get('w', shape=(2, 2))
def hybrid_forward(self, F, x, w):
return F.dot(x, w)
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
@np.use_np_array
class TestHybridBlock2(gluon.HybridBlock):
def __init__(self):
super(TestHybridBlock2, self).__init__()
self.w = self.params.get('w', shape=(2, 2))
def hybrid_forward(self, F, x, w):
return F.np.dot(x, w)
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):
@wraps_safely(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 `mxnet.numpy.ndarray` should be created instead of
`mx.nd.NDArray`.
Example::
import mxnet as mx
from mxnet import gluon, np
class TestHybridBlock1(gluon.HybridBlock):
def __init__(self):
super(TestHybridBlock1, self).__init__()
self.w = self.params.get('w', shape=(2, 2))
def hybrid_forward(self, F, x, w):
return F.dot(x, w) + F.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
@np.use_np
class TestHybridBlock2(gluon.HybridBlock):
def __init__(self):
super(TestHybridBlock2, self).__init__()
self.w = self.params.get('w', shape=(2, 2))
def hybrid_forward(self, F, x, w):
return F.np.dot(x, w) + F.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.
"""
@wraps_safely(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.
"""
@wraps_safely(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 _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.
"""
if active:
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.')
cur_state = is_np_array()
_NumpyArrayScope._current.value = _NumpyArrayScope(active)
return cur_state
def set_np(shape=True, array=True):
"""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`.
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.])]
"""
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)
def reset_np():
"""Deactivate NumPy shape and array semantics at the same time."""
set_np(shape=False, array=False)
_CUDA_SUCCESS = 0
def get_cuda_compute_capability(ctx):
"""Returns the cuda compute capability of the input `ctx`.
Parameters
----------
ctx : Context
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 ctx.device_type != 'gpu':
raise ValueError('Expecting a gpu context to get cuda compute capability, '
'while received ctx {}'.format(str(ctx)))
libnames = ('libcuda.so', 'libcuda.dylib', '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()
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(device), ctx.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), 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