我们从Python开源项目中，提取了以下50个代码示例，用于说明如何使用torch._C。
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") + ".")
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 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 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()
def _free_mutex(): torch._C._cuda_lock_mutex() try: yield finally: torch._C._cuda_unlock_mutex()
def once_differentiable(fn): from .variable import Variable @functools.wraps(fn) def wrapper(ctx, *args): tensor_args = [arg.data if isinstance(arg, Variable) else arg for arg in args] outputs = fn(ctx, *tensor_args) # XXX: this is only an approximation of these flags - there's no way # to figure out if fn didn't use ctx.saved_variables and as a result # some Variables might require grad, even if no args do. # Unfortunately, this leads to unexpected error messages ("no nodes # require computing gradients"), but I don't have a better idea. # These functions would raise an error in backward anyway. volatile = any(arg.volatile if isinstance(arg, Variable) else False for arg in args) requires_grad = any(arg.requires_grad if isinstance(arg, Variable) else False for arg in args) if volatile: def err_fn(*args): return args kwargs = {'volatile': True} else: err_fn = torch._C._functions.DelayedError( b"trying to differentiate twice a function that was marked" b"with @once_differentiable") kwargs = {'requires_grad': requires_grad} if not isinstance(outputs, tuple): var = Variable(outputs, **kwargs) if outputs is not None else None return err_fn(var) return err_fn(*[Variable(o, **kwargs) if o is not None else None for o in outputs]) return wrapper
def batch_norm(input, running_mean, running_var, weight=None, bias=None, training=False, momentum=0.1, eps=1e-5): f = torch._C._functions.BatchNorm(running_mean, running_var, training, momentum, eps, torch.backends.cudnn.enabled) return f(input, weight, bias) # loss
def _lazy_init(): global _initialized, _cudart, _original_pid, _queued_calls 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 for queued_call, orig_traceback in _queued_calls: try: queued_call() except Exception as e: msg = ("CUDA call failed lazily at initialization with error: {}\n\n" "CUDA call was originally invoked at:\n\n{}").format(str(e), orig_traceback) raise_from(DeferredCudaCallError(msg), e)
def __exit__(self, *args): if self.prev_idx != self.idx: torch._C._cuda_setDevice(self.prev_idx) return False
def get_device_name(device): """Gets the name of a device. Arguments: device (int): device for which to return the name. This function is a no-op if this argument is negative. """ if device >= 0: return torch._C._cuda_getDeviceName(device)
