Caffe2 - Python API
A deep learning, cross platform ML framework
__init__.py
1 import os
2 import ctypes
3 import sys
4 import torch
5 import types
6 import warnings
7 from torch.version import cuda
8 from contextlib import contextmanager
9 from subprocess import Popen, PIPE
10 
11 # Write:
12 #
13 # torch.backends.cudnn.enabled = False
14 #
15 # to globally disable CuDNN
16 
17 lib = None
18 __cudnn_version = None
19 # TODO: dynamic version checks via cudnnGetVersion
20 
21 
22 # The idea for this parameter is that we forbid bare assignment
23 # to torch.backends.cudnn.enabled and friends when running our
24 # test suite, where it's very easy to forget to undo the change
25 # later.
26 __allow_nonbracketed_mutation_flag = True
27 
28 
29 def find_cudnn_windows_lib():
30  proc = Popen(['where', 'cudnn64*.dll'], stdout=PIPE, stderr=PIPE, stdin=PIPE)
31  out, err = proc.communicate()
32  out = out.decode().strip()
33  if len(out) > 0:
34  if out.find('\r\n') != -1:
35  out = out.split('\r\n')[0]
36  cudnn_lib_name = os.path.basename(out)
37  cudnn_lib = os.path.splitext(cudnn_lib_name)[0]
38  cudnn_lib = str(cudnn_lib)
39  return ctypes.cdll.LoadLibrary(cudnn_lib)
40  else:
41  return None
42 
43 
44 def _libcudnn():
45  global lib, __cudnn_version
46  if lib is None:
47  if sys.platform == "win32":
48  lib = find_cudnn_windows_lib()
49  else:
50  lib = ctypes.cdll.LoadLibrary(None)
51  if hasattr(lib, 'cudnnGetErrorString'):
52  lib.cudnnGetErrorString.restype = ctypes.c_char_p
53  __cudnn_version = lib.cudnnGetVersion()
54  compile_version = torch._C._cudnn_version()
55  # cuDNN version is MAJOR*1000 + MINOR*100 + PATCH
56  runtime_major = __cudnn_version // 1000
57  runtime_minor = (__cudnn_version % 1000) // 100
58  compile_major = compile_version // 1000
59  compile_minor = (compile_version % 1000) // 100
60  # Different major versions are always incompatible
61  # Starting with cuDNN 7, minor versions are backwards-compatible
62  if runtime_major != compile_major:
63  cudnn_compatible = False
64  elif runtime_major < 7:
65  cudnn_compatible = runtime_minor == compile_minor
66  else:
67  cudnn_compatible = runtime_minor >= compile_minor
68  if not cudnn_compatible:
69  raise RuntimeError(
70  'cuDNN version incompatibility: PyTorch was compiled against {} '
71  'but linked against {}'.format(compile_version, __cudnn_version))
72  else:
73  lib = None
74  return lib
75 
76 
77 def version():
78  if _libcudnn() is None:
79  return None
80  return __cudnn_version
81 
82 
83 CUDNN_TENSOR_TYPES = {
84  'torch.cuda.HalfTensor',
85  'torch.cuda.FloatTensor',
86  'torch.cuda.DoubleTensor',
87 }
88 
89 
90 def is_available():
91  r"""Returns a bool indicating if CUDNN is currently available."""
92  return torch._C.has_cudnn
93 
94 
95 def is_acceptable(tensor):
96  if not torch._C._get_cudnn_enabled():
97  return False
98  if tensor.type() not in CUDNN_TENSOR_TYPES:
99  return False
100  if not is_available():
101  warnings.warn(
102  "PyTorch was compiled without cuDNN support. To use cuDNN, rebuild "
103  "PyTorch making sure the library is visible to the build system.")
104  return False
105  if _libcudnn() is None:
106  warnings.warn('cuDNN library not found. Check your {libpath}'.format(
107  libpath={
108  'darwin': 'DYLD_LIBRARY_PATH',
109  'win32': 'PATH'
110  }.get(sys.platform, 'LD_LIBRARY_PATH')))
111  return False
112  return True
113 
114 
115 _handles = {}
116 
117 verbose = False
118 
119 CUDNN_DATA_FLOAT = 0
120 CUDNN_DATA_DOUBLE = 1
121 CUDNN_DATA_HALF = 2
122 
123 CUDNN_TENSOR_NCHW = 0
124 CUDNN_TENSOR_NHWC = 1
125 
126 CUDNN_RNN_RELU = 0
127 CUDNN_RNN_TANH = 1
128 CUDNN_LSTM = 2
129 CUDNN_GRU = 3
130 
131 CUDNN_LINEAR_INPUT = 0
132 CUDNN_SKIP_INPUT = 1
133 
134 CUDNN_RNN_ALGO_STANDARD = 0
135 CUDNN_RNN_ALGO_PERSIST_STATIC = 1
136 CUDNN_RNN_ALGO_PERSIST_DYNAMIC = 2
137 
138 CUDNN_DEFAULT_MATH = 0
139 CUDNN_TENSOR_OP_MATH = 1
140 
141 
142 def set_flags(_enabled, _benchmark, _deterministic, _verbose):
143  global benchmark, deterministic, verbose
144  orig_flags = (torch._C._get_cudnn_enabled(),
145  torch._C._get_cudnn_benchmark(),
146  torch._C._get_cudnn_deterministic(),
147  verbose)
148  verbose = _verbose
149  torch._C._set_cudnn_enabled(_enabled)
150  torch._C._set_cudnn_benchmark(_benchmark)
151  torch._C._set_cudnn_deterministic(_deterministic)
152  return orig_flags
153 
154 
155 def disable_global_flags():
156  global __allow_nonbracketed_mutation_flag
157  __allow_nonbracketed_mutation_flag = False
158 
159 
160 def flags_frozen():
161  return not __allow_nonbracketed_mutation_flag
162 
163 
164 @contextmanager
165 def __allow_nonbracketed_mutation():
166  global __allow_nonbracketed_mutation_flag
167  old = __allow_nonbracketed_mutation_flag
168  __allow_nonbracketed_mutation_flag = True
169  try:
170  yield
171  finally:
172  __allow_nonbracketed_mutation_flag = old
173 
174 
175 @contextmanager
176 def flags(enabled=False, benchmark=False, deterministic=False, verbose=False):
177  with __allow_nonbracketed_mutation():
178  orig_flags = set_flags(enabled, benchmark, deterministic, verbose)
179  try:
180  yield
181  finally:
182  # recover the previous values
183  with __allow_nonbracketed_mutation():
184  set_flags(orig_flags[0], orig_flags[1], orig_flags[2], orig_flags[3])
185 
186 
188  def __init__(self):
189  ptr = ctypes.c_void_p()
190  check_error(lib.cudnnCreate(ctypes.byref(ptr)))
191  self._as_parameter_ = ptr
192 
193  def __del__(self):
194  check_error(lib.cudnnDestroy(self))
195 
196 
197 class CuDNNError(RuntimeError):
198  def __init__(self, status):
199  self.status = status
200  msg = '{}: {}'.format(status, get_error_string(status))
201  super(CuDNNError, self).__init__(msg)
202 
203 
204 class TensorDescriptor(object):
205  def __init__(self):
206  ptr = ctypes.c_void_p()
207  check_error(lib.cudnnCreateTensorDescriptor(ctypes.byref(ptr)))
208  self._as_parameter_ = ptr
209 
210  def __del__(self):
211  check_error(lib.cudnnDestroyTensorDescriptor(self._as_parameter_))
212  del self._as_parameter_
213 
214  def set(self, tensor):
215  self._type = tensor.type()
216  self._size = tensor.size()
217  self._stride = tensor.stride()
218  check_error(lib.cudnnSetTensorNdDescriptor(
219  self, _typemap[tensor.type()], tensor.dim(),
220  int_array(tensor.size()), int_array(tensor.stride())))
221 
222  def as_tuple(self):
223  return (self._type, tuple(self._size), tuple(self._stride))
224 
225 
226 class TensorDescriptorArray(object):
227  def __init__(self, N):
228  self.ptrs = (ctypes.c_void_p * N)()
229  for i in range(N):
230  ptr = ctypes.byref(self.ptrs, i * ctypes.sizeof(ctypes.c_void_p))
231  check_error(lib.cudnnCreateTensorDescriptor(ptr))
232  self._as_parameter_ = self.ptrs
233 
234  def __del__(self):
235  for ptr in self.ptrs:
236  check_error(lib.cudnnDestroyTensorDescriptor(ctypes.c_void_p(ptr)))
237 
238  def __getitem__(self, key):
239  return ctypes.c_void_p(self.ptrs[key])
240 
241  def set_all(self, tensor):
242  _type = _typemap[tensor.type()]
243  _ndim = tensor.dim()
244  _size = int_array(tensor.size())
245  _stride = int_array(tensor.stride())
246  for ptr in self.ptrs:
247  check_error(lib.cudnnSetTensorNdDescriptor(
248  ctypes.c_void_p(ptr), _type, _ndim, _size, _stride))
249 
250  def set_raw(self, i, _type, _ndim, _size, _stride):
251  ptr = self.ptrs[i]
252  check_error(lib.cudnnSetTensorNdDescriptor(
253  ctypes.c_void_p(ptr), _type, _ndim, _size, _stride))
254 
255 
256 class FilterDescriptor(object):
257  def __init__(self):
258  ptr = ctypes.c_void_p()
259  check_error(lib.cudnnCreateFilterDescriptor(ctypes.byref(ptr)))
260  self._as_parameter_ = ptr
261 
262  def __del__(self):
263  check_error(lib.cudnnDestroyFilterDescriptor(self._as_parameter_))
264  del self._as_parameter_
265 
266  def set(self, weight):
267  self._size = weight.size()
268  datatype = _typemap[weight.type()]
269  check_error(lib.cudnnSetFilterNdDescriptor(
270  self, datatype, CUDNN_TENSOR_NCHW, weight.ndimension(),
271  int_array(weight.size())))
272 
273  def as_tuple(self):
274  return tuple(self._size)
275 
276 
277 class DropoutDescriptor(object):
278  def __init__(self, handle, dropout, seed):
279  ptr = ctypes.c_void_p()
280  check_error(lib.cudnnCreateDropoutDescriptor(ctypes.byref(ptr)))
281 
282  self._as_parameter_ = ptr
283  self.state = None
284  self.dropout = dropout
285  self.handle = handle
286 
287  self._set(dropout, seed)
288 
289  def set_dropout(self, dropout, seed):
290  if dropout != self.dropout:
291  self._set(dropout, seed)
292 
293  def _set(self, dropout, seed):
294  if self.state is None and dropout > 0:
295  dropout_states_size = ctypes.c_long()
296  check_error(lib.cudnnDropoutGetStatesSize(
297  self.handle,
298  ctypes.byref(dropout_states_size)))
299  self.state = torch.cuda.ByteTensor(dropout_states_size.value)
300  state_ptr = self.state.data_ptr()
301  state_size = self.state.size(0)
302  else:
303  state_ptr = None
304  state_size = 0
305 
306  check_error(lib.cudnnSetDropoutDescriptor(
307  self,
308  self.handle,
309  ctypes.c_float(dropout),
310  ctypes.c_void_p(state_ptr),
311  ctypes.c_size_t(state_size),
312  ctypes.c_ulonglong(seed),
313  ))
314 
315  self.dropout = dropout
316 
317  def __del__(self):
318  check_error(lib.cudnnDestroyDropoutDescriptor(self))
319 
320 
321 class RNNDescriptor(object):
322  def __init__(self, handle, hidden_size, num_layers, dropout_desc, input_mode,
323  bidirectional, mode, datatype):
324  ptr = ctypes.c_void_p()
325  check_error(lib.cudnnCreateRNNDescriptor(ctypes.byref(ptr)))
326  self._as_parameter_ = ptr
327  if version() >= 6000:
328  check_error(lib.cudnnSetRNNDescriptor_v6(
329  handle,
330  self,
331  hidden_size,
332  num_layers,
333  dropout_desc,
334  input_mode,
335  bidirectional,
336  mode,
337  CUDNN_RNN_ALGO_STANDARD,
338  datatype
339  ))
340  if version() >= 7000 and int(cuda[0]) >= 9 and (
342  lib.cudnnSetRNNMatrixMathType(self, CUDNN_DEFAULT_MATH)
343  if datatype == CUDNN_DATA_HALF:
344  lib.cudnnSetRNNMatrixMathType(self, CUDNN_TENSOR_OP_MATH)
345  else:
346  check_error(lib.cudnnSetRNNDescriptor(
347  self,
348  hidden_size,
349  num_layers,
350  dropout_desc,
351  input_mode,
352  bidirectional,
353  mode,
354  datatype
355  ))
356 
357  def __del__(self):
358  check_error(lib.cudnnDestroyRNNDescriptor(self))
359 
360 
361 def check_error(status):
362  if status != 0:
363  raise CuDNNError(status)
364 
365 
366 def get_error_string(status):
367  return lib.cudnnGetErrorString(status)
368 
369 
370 def get_handle():
371  if _libcudnn() is None:
372  raise RuntimeError('cuDNN not available')
373  current_device = torch.cuda.current_device()
374  handle = _handles.get(current_device, None)
375  if handle is None:
376  handle = CuDNNHandle()
377  _handles[current_device] = handle
378  return handle
379 
380 
381 _typemap = {
382  'torch.cuda.HalfTensor': CUDNN_DATA_HALF,
383  'torch.cuda.FloatTensor': CUDNN_DATA_FLOAT,
384  'torch.cuda.DoubleTensor': CUDNN_DATA_DOUBLE,
385 }
386 
387 _sizeofmap = {
388  CUDNN_DATA_HALF: 2,
389  CUDNN_DATA_FLOAT: 4,
390  CUDNN_DATA_DOUBLE: 8,
391 }
392 
393 
394 def c_type(tensor):
395  if isinstance(tensor, torch.cuda.HalfTensor):
396  return ctypes.c_float
397  elif isinstance(tensor, torch.cuda.FloatTensor):
398  return ctypes.c_float
399  elif isinstance(tensor, torch.cuda.DoubleTensor):
400  return ctypes.c_double
401  else:
402  raise ValueError("unknown type '{}'".format(type(tensor)))
403 
404 
405 def int_array(itr):
406  array_type = ctypes.c_int * len(itr)
407  return array_type(*itr)
408 
409 
410 def descriptor(tensor, N=None):
411  padded_size = tensor.size() + ((1,) * (5 - tensor.dim()))
412  tensor = tensor.view(padded_size)
413  if N is not None:
414  descriptor = TensorDescriptorArray(N)
415  descriptor.set_all(tensor)
416  else:
417  descriptor = TensorDescriptor()
418  descriptor.set(tensor)
419  return descriptor
420 
421 
422 def descriptor_sequence(tensor, batch_sizes):
423  descriptors = TensorDescriptorArray(len(batch_sizes))
424  _type = _typemap[tensor.type()]
425  _ndim = 5
426  dim_pad = (1,) * (5 - tensor.dim())
427  _size = int_array(tensor.size() + dim_pad)
428  _stride = int_array(tensor.stride() + dim_pad)
429  for i, batch_size in enumerate(batch_sizes):
430  _size[0] = batch_size
431  descriptors.set_raw(i, _type, _ndim, _size, _stride)
432  return descriptors
433 
434 
435 def add_tensor(*args):
436  check_error(lib.cudnnAddTensor(*args))
437 
438 
439 # The magic here is to allow us to intercept code like this:
440 #
441 # torch.backends.cudnn.enabled = True
442 
443 class ContextProp(object):
444  def __init__(self, getter, setter):
445  self.getter = getter
446  self.setter = setter
447 
448  def __get__(self, obj, objtype):
449  return self.getter()
450 
451  def __set__(self, obj, val):
452  if not flags_frozen():
453  self.setter(val)
454  else:
455  raise RuntimeError("not allowed to set torch.backends.cudnn flags "
456  "after disable_global_flags; please use flags() context manager instead")
457 
458 
459 class CudnnModule(types.ModuleType):
460  def __init__(self, m, name):
461  super(CudnnModule, self).__init__(name)
462  self.m = m
463 
464  def __getattr__(self, attr):
465  return self.m.__getattribute__(attr)
466 
467  enabled = ContextProp(torch._C._get_cudnn_enabled, torch._C._set_cudnn_enabled)
468  deterministic = ContextProp(torch._C._get_cudnn_deterministic, torch._C._set_cudnn_deterministic)
469  benchmark = ContextProp(torch._C._get_cudnn_benchmark, torch._C._set_cudnn_benchmark)
470 
471 # This is the sys.modules replacement trick, see
472 # https://stackoverflow.com/questions/2447353/getattr-on-a-module/7668273#7668273
473 sys.modules[__name__] = CudnnModule(sys.modules[__name__], __name__)
def _set(self, dropout, seed)
Definition: __init__.py:293
def get_device_capability(device=None)
Definition: __init__.py:280
def current_device()
Definition: __init__.py:349