def _load_cudart():
system = platform.system()
lib_name = 'libcudart.' + ('dylib' if system == 'Darwin' else 'so')
lib_paths = [
lib_name,
os.path.join(torch._C._cuda_getLibPath(), lib_name),
os.path.join('/usr/local/cuda/lib64', lib_name),
os.path.join('/usr/local/cuda/lib', lib_name),
]
for path in lib_paths:
try:
return ctypes.cdll.LoadLibrary(path)
except OSError:
pass
raise RuntimeError("couldn't find libcudart. Make sure CUDA libraries "
"are installed in a default location, or that they're in " +
("DYLD_LIBRARY_PATH" if system == 'Darwin' else "LD_LIBRARY_PATH") +
".")
python类_C的实例源码
def _check_driver():
if not hasattr(torch._C, '_cuda_isDriverSufficient'):
raise AssertionError("Torch not compiled with CUDA enabled")
if not torch._C._cuda_isDriverSufficient():
if torch._C._cuda_getDriverVersion() == 0:
# found no NVIDIA driver on the system
raise AssertionError("""
Found no NVIDIA driver on your system. Please check that you
have an NVIDIA GPU and installed a driver from
http://www.nvidia.com/Download/index.aspx""")
else:
# TODO: directly link to the alternative bin that needs install
raise AssertionError("""
The NVIDIA driver on your system is too old (found version {}).
Please update your GPU driver by downloading and installing a new
version from the URL: http://www.nvidia.com/Download/index.aspx
Alternatively, go to: https://pytorch.org/binaries to install
a PyTorch version that has been compiled with your version
of the CUDA driver.""".format(str(torch._C._cuda_getDriverVersion())))
def _check_driver():
if not hasattr(torch._C, '_cuda_isDriverSufficient'):
raise AssertionError("Torch not compiled with CUDA enabled")
if not torch._C._cuda_isDriverSufficient():
if torch._C._cuda_getDriverVersion() == 0:
# found no NVIDIA driver on the system
raise AssertionError("""
Found no NVIDIA driver on your system. Please check that you
have an NVIDIA GPU and installed a driver from
http://www.nvidia.com/Download/index.aspx""")
else:
# TODO: directly link to the alternative bin that needs install
raise AssertionError("""
The NVIDIA driver on your system is too old (found version {}).
Please update your GPU driver by downloading and installing a new
version from the URL: http://www.nvidia.com/Download/index.aspx
Alternatively, go to: https://pytorch.org/binaries to install
a PyTorch version that has been compiled with your version
of the CUDA driver.""".format(str(torch._C._cuda_getDriverVersion())))
def _lazy_init():
global _initialized, _cudart, _original_pid
if _initialized:
return
if _in_bad_fork:
from sys import version_info
if version_info < (3, 4):
msg = ("To use CUDA with multiprocessing, you must use Python "
"3.4+ and the 'spawn' start method")
else:
msg = ("To use CUDA with multiprocessing, you must use the "
"'spawn' start method")
raise RuntimeError(
"Cannot re-initialize CUDA in forked subprocess. " + msg)
_check_driver()
torch._C._cuda_init()
torch._C._cuda_sparse_init()
_cudart = _load_cudart()
_cudart.cudaGetErrorName.restype = ctypes.c_char_p
_cudart.cudaGetErrorString.restype = ctypes.c_char_p
_original_pid = os.getpid()
_initialized = True
def stream(stream):
"""Context-manager that selects a given stream.
All CUDA kernels queued within its context will be enqueued on a selected
stream.
Arguments:
stream (Stream): selected stream. This manager is a no-op if it's
``None``.
"""
if stream is None:
yield
return
prev_stream = current_stream()
torch._C._cuda_setStream(stream._cdata)
try:
yield
finally:
torch._C._cuda_setStream(prev_stream._cdata)
def _check_driver():
if not hasattr(torch._C, '_cuda_isDriverSufficient'):
raise AssertionError("Torch not compiled with CUDA enabled")
if not torch._C._cuda_isDriverSufficient():
if torch._C._cuda_getDriverVersion() == 0:
# found no NVIDIA driver on the system
raise AssertionError("""
Found no NVIDIA driver on your system. Please check that you
have an NVIDIA GPU and installed a driver from
http://www.nvidia.com/Download/index.aspx""")
else:
# TODO: directly link to the alternative bin that needs install
raise AssertionError("""
The NVIDIA driver on your system is too old (found version {}).
Please update your GPU driver by downloading and installing a new
version from the URL: http://www.nvidia.com/Download/index.aspx
Alternatively, go to: https://pytorch.org/binaries to install
a PyTorch version that has been compiled with your version
of the CUDA driver.""".format(str(torch._C._cuda_getDriverVersion())))
def _lazy_init():
global _initialized, _cudart, _original_pid
if _initialized:
return
if _in_bad_fork:
from sys import version_info
if version_info < (3, 4):
msg = ("To use CUDA with multiprocessing, you must use Python "
"3.4+ and the 'spawn' start method")
else:
msg = ("To use CUDA with multiprocessing, you must use the "
"'spawn' start method")
raise RuntimeError(
"Cannot re-initialize CUDA in forked subprocess. " + msg)
_check_driver()
torch._C._cuda_init()
torch._C._cuda_sparse_init()
_cudart = _load_cudart()
_cudart.cudaGetErrorName.restype = ctypes.c_char_p
_cudart.cudaGetErrorString.restype = ctypes.c_char_p
_original_pid = os.getpid()
_initialized = True
def stream(stream):
"""Context-manager that selects a given stream.
All CUDA kernels queued within its context will be enqueued on a selected
stream.
Arguments:
stream (Stream): selected stream. This manager is a no-op if it's
``None``.
"""
if stream is None:
yield
return
prev_stream = current_stream()
torch._C._cuda_setStream(stream._cdata)
try:
yield
finally:
torch._C._cuda_setStream(prev_stream._cdata)
def _check_driver():
if not hasattr(torch._C, '_cuda_isDriverSufficient'):
raise AssertionError("Torch not compiled with CUDA enabled")
if not torch._C._cuda_isDriverSufficient():
if torch._C._cuda_getDriverVersion() == 0:
# found no NVIDIA driver on the system
raise AssertionError("""
Found no NVIDIA driver on your system. Please check that you
have an NVIDIA GPU and installed a driver from
http://www.nvidia.com/Download/index.aspx""")
else:
# TODO: directly link to the alternative bin that needs install
raise AssertionError("""
The NVIDIA driver on your system is too old (found version {}).
Please update your GPU driver by downloading and installing a new
version from the URL: http://www.nvidia.com/Download/index.aspx
Alternatively, go to: https://pytorch.org/binaries to install
a PyTorch version that has been compiled with your version
of the CUDA driver.""".format(str(torch._C._cuda_getDriverVersion())))
def stream(stream):
"""Context-manager that selects a given stream.
All CUDA kernels queued within its context will be enqueued on a selected
stream.
Arguments:
stream (Stream): selected stream. This manager is a no-op if it's
``None``.
"""
if stream is None:
yield
return
prev_stream = current_stream()
torch._C._cuda_setStream(stream._cdata)
try:
yield
finally:
torch._C._cuda_setStream(prev_stream._cdata)
def log_softmax(input, dim=None, _stacklevel=3):
r"""Applies a softmax followed by a logarithm.
While mathematically equivalent to log(softmax(x)), doing these two
operations separately is slower, and numerically unstable. This function
uses an alternative formulation to compute the output and gradient correctly.
See :class:`~torch.nn.LogSoftmax` for more details.
Arguments:
input (Variable): input
dim (int): A dimension along which log_softmax will be computed.
"""
if dim is None:
dim = _get_softmax_dim('log_softmax', input.dim(), _stacklevel)
return torch._C._nn.log_softmax(input, dim)
def _check_driver():
if not hasattr(torch._C, '_cuda_isDriverSufficient'):
raise AssertionError("Torch not compiled with CUDA enabled")
if not torch._C._cuda_isDriverSufficient():
if torch._C._cuda_getDriverVersion() == 0:
# found no NVIDIA driver on the system
raise AssertionError("""
Found no NVIDIA driver on your system. Please check that you
have an NVIDIA GPU and installed a driver from
http://www.nvidia.com/Download/index.aspx""")
else:
# TODO: directly link to the alternative bin that needs install
raise AssertionError("""
The NVIDIA driver on your system is too old (found version {}).
Please update your GPU driver by downloading and installing a new
version from the URL: http://www.nvidia.com/Download/index.aspx
Alternatively, go to: http://pytorch.org to install
a PyTorch version that has been compiled with your version
of the CUDA driver.""".format(str(torch._C._cuda_getDriverVersion())))
def _check_capability():
error_str = """
Found GPU%d %s which requires CUDA_VERSION >= %d for
optimal performance and fast startup time, but your PyTorch was compiled
with CUDA_VERSION %d. Please install the correct PyTorch binary
using instructions from http://pytorch.org
"""
CUDA_VERSION = torch._C._cuda_getCompiledVersion()
for d in range(device_count()):
major = get_device_capability(d)[0]
name = get_device_name(d)
if CUDA_VERSION < 8000 and major >= 6:
warnings.warn(error_str % (d, name, 8000, CUDA_VERSION))
elif CUDA_VERSION < 9000 and major >= 7:
warnings.warn(error_str % (d, name, 9000, CUDA_VERSION))
def is_available():
return (hasattr(torch._C, '_cuda_isDriverSufficient') and
torch._C._cuda_isDriverSufficient())
def _lazy_init():
global _initialized, _cudart
if _initialized:
return
_check_driver()
assert torch._C._cuda_init()
_cudart = _load_cudart()
_cudart.cudaGetErrorName.restype = ctypes.c_char_p
_cudart.cudaGetErrorString.restype = ctypes.c_char_p
_initialized = True
def __enter__(self):
if self.idx is -1:
return
_lazy_init()
self.prev_idx = torch._C._cuda_getDevice()
if self.prev_idx != self.idx:
torch._C._cuda_setDevice(self.idx)
def stream(stream):
if stream is None:
yield
return
prev_stream = current_stream()
torch._C._cuda_setStream(stream._cdata)
try:
yield
finally:
torch._C._cuda_setStream(prev_stream._cdata)
def device_count():
_lazy_init()
return torch._C._cuda_getDeviceCount()
def current_device():
_lazy_init()
return torch._C._cuda_getDevice()
def synchronize():
_lazy_init()
return torch._C._cuda_synchronize()
def current_stream():
_lazy_init()
return torch.cuda.Stream(_cdata=torch._C._cuda_getCurrentStream())
def is_available():
"""Returns a bool indicating if CUDA is currently available."""
if (not hasattr(torch._C, '_cuda_isDriverSufficient') or
not torch._C._cuda_isDriverSufficient()):
return False
return torch._C._cuda_getDeviceCount() > 0
def _sleep(cycles):
torch._C._cuda_sleep(cycles)
def __enter__(self):
if self.idx is -1:
return
_lazy_init()
self.prev_idx = torch._C._cuda_getDevice()
if self.prev_idx != self.idx:
torch._C._cuda_setDevice(self.idx)
def set_device(device):
"""Sets the current device.
Usage of this function is discouraged in favor of :any:`device`. In most
cases it's better to use ``CUDA_VISIBLE_DEVICES`` environmental variable.
Arguments:
device (int): selected device. This function is a no-op if this
argument is negative.
"""
if device >= 0:
torch._C._cuda_setDevice(device)
def device_count():
"""Returns the number of GPUs available."""
if is_available():
_lazy_init()
return torch._C._cuda_getDeviceCount()
else:
return 0
def current_device():
"""Returns the index of a currently selected device."""
_lazy_init()
return torch._C._cuda_getDevice()
def synchronize():
"""Waits for all kernels in all streams on current device to complete."""
_lazy_init()
return torch._C._cuda_synchronize()
def current_blas_handle():
"""Returns cublasHandle_t pointer to current cuBLAS handle"""
return torch._C._cuda_getCurrentBlasHandle()
def _host_allocator():
_lazy_init()
return torch._C._cuda_cudaHostAllocator()
