def _update_output(self, input, weight, bias):
self.use_cudnn = cudnn.is_acceptable(input)
if self.use_cudnn and cudnn.version() < 6000:
self.use_cudnn = not self.is_dilated()
if self.use_cudnn:
output = input.new(*self._output_size(input, weight))
if self.transposed:
self._cudnn_info = (
torch._C._cudnn_convolution_transpose_full_forward(
input, weight, bias, output, self.padding, self.stride, self.dilation,
self.groups, cudnn.benchmark))
else:
self._cudnn_info = torch._C._cudnn_convolution_full_forward(
input, weight, bias, output, self.padding, self.stride, self.dilation,
self.groups, cudnn.benchmark)
if not self.requires_grad:
del self._cudnn_info
return output
self._bufs = [[] for g in range(self.groups)]
output = self._thnn('update_output', input, weight, bias)
if not self.requires_grad:
del self._bufs
return output
python类_thnn()的实例源码
def _grad_input(self, input, weight, grad_output):
if self.use_cudnn:
grad_input = input.new().resize_as_(input)
if self.transposed:
# ConvTranspose uses the same kernels as regular convolution
# but swaps forward and backward calls
torch._C._cudnn_convolution_forward(
grad_output, weight, grad_input, self._cudnn_info,
cudnn.benchmark)
else:
torch._C._cudnn_convolution_backward_data(
grad_output, grad_input, weight, self._cudnn_info,
cudnn.benchmark)
return grad_input
return self._thnn('grad_input', input, weight, grad_output)
def _grad_params(self, input, weight, bias, grad_output):
if self.use_cudnn:
grad_weight = grad_bias = None
if self.needs_input_grad[1]:
grad_weight = weight.new().resize_as_(weight)
torch._C._cudnn_convolution_backward_filter(
grad_output, input, grad_weight, self._cudnn_info,
cudnn.benchmark)
if bias is not None and self.needs_input_grad[2]:
grad_bias = bias.new().resize_as_(bias)
torch._C._cudnn_convolution_backward_bias(
grad_output, grad_bias, self._cudnn_info)
return grad_weight, grad_bias
return self._thnn('grad_params', input, weight, bias, grad_output)
def _thnn(self, fn_name, input, weight, *args):
impl = _thnn_convs[self.thnn_class_name(input)]
if self.groups == 1:
return impl[fn_name](self, self._bufs[0], input, weight, *args)
else:
res = []
for g in range(self.groups):
def group(tensor, dim=None):
if tensor is None:
return None
if dim is None:
dim = 0 if tensor.dim() == 1 else 1
n = tensor.size(dim) // self.groups
return tensor.narrow(dim, n * g, n).contiguous()
grouped_args = [group(input, 1), group(weight, 0)]
grouped_args += [group(t) for t in args]
res.append(impl[fn_name](self, self._bufs[g], *grouped_args))
if fn_name == 'grad_params':
return [torch.cat(t, 0) if t[0] is not None else None
for t in zip(*res)]
else:
return torch.cat(res, 1)
def _update_output(self, input, weight, bias):
self.use_cudnn = cudnn.is_acceptable(input)
if self.use_cudnn and cudnn.version() < 6000:
self.use_cudnn = not self.is_dilated()
if self.use_cudnn:
output = input.new(*self._output_size(input, weight))
if self.transposed:
self._cudnn_info = (
torch._C._cudnn_convolution_transpose_full_forward(
input, weight, bias, output, self.padding, self.stride, self.dilation,
self.groups, cudnn.benchmark))
else:
self._cudnn_info = torch._C._cudnn_convolution_full_forward(
input, weight, bias, output, self.padding, self.stride, self.dilation,
self.groups, cudnn.benchmark)
if not self.requires_grad:
del self._cudnn_info
return output
self._bufs = [[] for g in range(self.groups)]
output = self._thnn('update_output', input, weight, bias)
if not self.requires_grad:
del self._bufs
return output
def _grad_input(self, input, weight, grad_output):
if self.use_cudnn:
grad_input = input.new().resize_as_(input)
if self.transposed:
# ConvTranspose uses the same kernels as regular convolution
# but swaps forward and backward calls
torch._C._cudnn_convolution_forward(
grad_output, weight, grad_input, self._cudnn_info,
cudnn.benchmark)
else:
torch._C._cudnn_convolution_backward_data(
grad_output, grad_input, weight, self._cudnn_info,
cudnn.benchmark)
return grad_input
return self._thnn('grad_input', input, weight, grad_output)
def _grad_params(self, input, weight, bias, grad_output):
if self.use_cudnn:
grad_weight = grad_bias = None
if self.needs_input_grad[1]:
grad_weight = weight.new().resize_as_(weight)
torch._C._cudnn_convolution_backward_filter(
grad_output, input, grad_weight, self._cudnn_info,
cudnn.benchmark)
if bias is not None and self.needs_input_grad[2]:
grad_bias = bias.new().resize_as_(bias)
torch._C._cudnn_convolution_backward_bias(
grad_output, grad_bias, self._cudnn_info)
return grad_weight, grad_bias
return self._thnn('grad_params', input, weight, bias, grad_output)
def _thnn(self, fn_name, input, weight, *args):
impl = _thnn_convs[self.thnn_class_name(input)]
if self.groups == 1:
return impl[fn_name](self, self._bufs[0], input, weight, *args)
else:
res = []
for g in range(self.groups):
def group(tensor, dim=None):
if tensor is None:
return None
if dim is None:
dim = 0 if tensor.dim() == 1 else 1
n = tensor.size(dim) // self.groups
return tensor.narrow(dim, n * g, n).contiguous()
grouped_args = [group(input, 1), group(weight, 0)]
grouped_args += [group(t) for t in args]
res.append(impl[fn_name](self, self._bufs[g], *grouped_args))
if fn_name == 'grad_params':
return [torch.cat(t, 0) if t[0] is not None else None
for t in zip(*res)]
else:
return torch.cat(res, 1)
def __init__(self):
self.gradInput = torch.Tensor()
self.output = 0
self._backend = torch._thnn.type2backend[type(self.gradInput)]
def type(self, type, tensorCache=None):
# find all tensors and convert them
for key, param in self.__dict__.items():
setattr(self, key, recursiveType(param, type, tensorCache or {}))
self._backend = torch._thnn.type2backend[type]
return self
def __init__(self):
self.gradInput = torch.Tensor()
self.output = torch.Tensor()
self._type = self.output.type()
self._backend = torch._thnn.type2backend[type(self.output)]
def type(self, type=None, tensorCache=None):
if not type:
return self._type
tensorCache = tensorCache or {}
# find all tensors and convert them
for key, param in self.__dict__.items():
setattr(self, key, recursiveType(param, type, tensorCache))
self._backend = torch._thnn.type2backend[type]
self._type = type
return self
def __init__(self):
self.gradInput = torch.Tensor()
self.output = 0
self._backend = torch._thnn.type2backend[type(self.gradInput)]
def __init__(self):
self.gradInput = torch.Tensor()
self.output = torch.Tensor()
self._type = self.output.type()
self._backend = torch._thnn.type2backend[type(self.output)]
def type(self, type=None, tensorCache=None):
if type is None:
return self._type
tensorCache = tensorCache or {}
# find all tensors and convert them
for key, param in self.__dict__.items():
setattr(self, key, recursiveType(param, type, tensorCache))
self._backend = torch._thnn.type2backend[type]
self._type = type
return self
def __init__(self):
self.gradInput = torch.Tensor()
self.output = 0
self._backend = torch._thnn.type2backend[type(self.gradInput)]
def __init__(self):
self.gradInput = torch.Tensor()
self.output = torch.Tensor()
self._type = self.output.type()
self._backend = torch._thnn.type2backend[type(self.output)]
def type(self, type=None, tensorCache=None):
if type is None:
return self._type
tensorCache = tensorCache or {}
# find all tensors and convert them
for key, param in self.__dict__.items():
setattr(self, key, recursiveType(param, type, tensorCache))
self._backend = torch._thnn.type2backend[type]
self._type = type
return self
def __init__(self):
self.gradInput = torch.Tensor()
self.output = 0
self._backend = torch._thnn.type2backend[type(self.gradInput)]
def type(self, type, tensorCache=None):
# find all tensors and convert them
for key, param in self.__dict__.items():
setattr(self, key, recursiveType(param, type, tensorCache or {}))
self._backend = torch._thnn.type2backend[type]
return self
def __init__(self):
self.gradInput = torch.Tensor()
self.output = torch.Tensor()
self._type = self.output.type()
self._backend = torch._thnn.type2backend[type(self.output)]
def type(self, type=None, tensorCache=None):
if type is None:
return self._type
tensorCache = tensorCache or {}
# find all tensors and convert them
for key, param in self.__dict__.items():
setattr(self, key, recursiveType(param, type, tensorCache))
self._backend = torch._thnn.type2backend[type]
self._type = type
return self
def __init__(self):
self.gradInput = torch.Tensor()
self.output = 0
self._backend = torch._thnn.type2backend[type(self.gradInput)]
def type(self, type, tensorCache=None):
# find all tensors and convert them
for key, param in self.__dict__.items():
setattr(self, key, recursiveType(param, type, tensorCache or {}))
self._backend = torch._thnn.type2backend[type]
return self
def __init__(self):
self.gradInput = torch.Tensor()
self.output = torch.Tensor()
self._type = self.output.type()
self._backend = torch._thnn.type2backend[type(self.output)]
def type(self, type=None, tensorCache=None):
if type is None:
return self._type
tensorCache = tensorCache or {}
# find all tensors and convert them
for key, param in self.__dict__.items():
setattr(self, key, recursiveType(param, type, tensorCache))
self._backend = torch._thnn.type2backend[type]
self._type = type
return self
