def _enforce_cudnn(input):
if not cudnn.enabled:
raise RuntimeError("AffineGridGenerator needs CuDNN for "
"processing CUDA inputs, but CuDNN is not enabled")
assert cudnn.is_acceptable(input)
python类enabled()的实例源码
def use_cudnn(should_use):
orig = torch.backends.cudnn.enabled
torch.backends.cudnn.enabled = should_use
try:
yield
finally:
torch.backends.cudnn.enabled = orig
def test_Conv2d_inconsistent_types_on_GPU_without_cudnn(self):
inputs = Variable(torch.randn(4, 1, 7, 7).float().cuda())
weights = Variable(torch.randn(1, 1, 3, 3).double().cuda())
bias = Variable(torch.randn(1).double().cuda())
torch.backends.cudnn.enabled = False
# inconsistent types should raise an exception
self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))
# but it should work with the same type
nn.functional.conv2d(inputs.float(), weights.float(), bias.float())
def test_Conv2d_inconsistent_types_on_GPU_with_cudnn(self):
inputs = Variable(torch.randn(4, 1, 7, 7).float().cuda())
weights = Variable(torch.randn(1, 1, 3, 3).double().cuda())
bias = Variable(torch.randn(1).double().cuda())
torch.backends.cudnn.enabled = True
# inconsistent types should raise an exception
self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))
# but it should work with the same type
nn.functional.conv2d(inputs.float(), weights.float(), bias.float())
def test_rnn_retain_variables_cuda(self):
try:
torch.backends.cudnn.enabled = False
self._test_rnn_retain_variables(torch.cuda.FloatTensor)
finally:
torch.backends.cudnn.enabled = True
self._test_rnn_retain_variables(torch.cuda.FloatTensor)
def affine_grid_generator(theta, size):
if theta.data.is_cuda:
if not cudnn.enabled:
raise RuntimeError("AffineGridGenerator needs CuDNN for "
"processing CUDA inputs, but CuDNN is not enabled")
if not cudnn.is_acceptable(theta.data):
raise RuntimeError("AffineGridGenerator generator theta not acceptable for CuDNN")
N, C, H, W = size
return torch._C._VariableBase.cudnn_affine_grid_generator(theta, N, C, H, W)
else:
return AffineGridGenerator.apply(theta, size)
# TODO: Port these completely into C++
def _enforce_cudnn(input):
if not cudnn.enabled:
raise RuntimeError("AffineGridGenerator needs CuDNN for "
"processing CUDA inputs, but CuDNN is not enabled")
assert cudnn.is_acceptable(input)