dnn应用的python代码_Python cudnn.lib方法代码示例

# 需要导入模块: from torch.backends import cudnn [as 别名]

# 或者: from torch.backends.cudnn import lib [as 别名]

def backward_weight(fn, input, hx, output, weight, grad_weight):

with torch.cuda.device_of(input):

is_input_packed = fn.batch_sizes is not None

handle = cudnn.get_handle()

if fn.mode == cudnn.CUDNN_LSTM:

hx, cx = hx

else:

cx = None

if fn.batch_first and not is_input_packed:

input = input.transpose(0, 1)

output = output.transpose(0, 1)

input_size = _input_size(fn, input)

hidden_size = _hidden_size(fn)

if not fn.requires_grad:

raise RuntimeError('backward_weight can only be called when the function requires grad!')

if fn.dropout != 0 and cudnn.version() < 5103:

raise RuntimeError('dropout supported only in cudnn v 5.1 and above')