Implement MaxPool{1d,2d,3d}Backwards (non-differentiable) functions.

Author: gchanan

torch/nn/_functions/thnn/pooling.py

@@ -43,14 +43,28 @@ def forward(ctx, input, kernel_size, stride=None, padding=0, dilation=1,
  return output
 
  @staticmethod
- @once_differentiable
  def backward(ctx, grad_output, _indices_grad=None):
  if ctx.return_indices:
- input, indices = ctx.saved_tensors
+ input, indices = ctx.saved_variables
  else:
- input, = ctx.saved_tensors
+ input, = ctx.saved_variables
  indices = ctx.indices
 
+ grad_input = MaxPool1dBackward.apply(input, indices, grad_output, ctx.kernel_size, ctx.stride, ctx.pad,
+ ctx.dilation, ctx.return_indices, ctx.ceil_mode)
+ return grad_input, None, None, None, None, None, None
+
+
+class MaxPool1dBackward(Function):
+
+ @staticmethod
+ def forward(ctx, input, indices, grad_output, kernel_size, stride, padding, dilation, return_indices, ceil_mode):
+ ctx.kernel_size = kernel_size
+ ctx.stride = stride
+ ctx.pad = padding
+ ctx.dilation = dilation
+ ctx.return_indices = return_indices
+ ctx.ceil_mode = ceil_mode
  input2d = input.unsqueeze(2)
  indices2d = indices.unsqueeze(2)
  grad_output2d = grad_output.unsqueeze(2)
@@ -64,7 +78,11 @@ def backward(ctx, grad_output, _indices_grad=None):
  ctx.dilation, 1,
  ctx.ceil_mode)
  grad_input = grad_input.squeeze(2)
- return grad_input, None, None, None, None, None, None
+ return grad_input
+
+ @staticmethod
+ def backward(ctx, ggI, ggIndices=None):
+ raise ValueError("MaxPool1d cannot be differentiated twice")
 
 
 class MaxPool2d(Function):
@@ -97,13 +115,29 @@ def forward(ctx, input, kernel_size, stride=None, padding=0, dilation=1,
  return output
 
  @staticmethod
- @once_differentiable
  def backward(ctx, grad_output, _indices_grad=None):
  if ctx.return_indices:
- input, indices = ctx.saved_tensors
+ input, indices = ctx.saved_variables
  else:
- input, = ctx.saved_tensors
+ input, = ctx.saved_variables
  indices = ctx.indices
+ grad_input = MaxPool2dBackward.apply(input, indices, grad_output, ctx.kernel_size, ctx.stride, ctx.padding,
+ ctx.dilation, ctx.return_indices, ctx.ceil_mode)
+ return grad_input, None, None, None, None, None, None
+
+
+class MaxPool2dBackward(Function):
+
+ @staticmethod
+ def forward(ctx, input, indices, grad_output, kernel_size, stride, padding, dilation,
+ return_indices, ceil_mode):
+ ctx.kernel_size = kernel_size
+ ctx.stride = stride
+ ctx.padding = padding
+ ctx.dilation = dilation
+ ctx.return_indices = return_indices
+ ctx.ceil_mode = ceil_mode
+
  grad_input = grad_output.new()
  backend = type2backend[type(input)]
  backend.SpatialDilatedMaxPooling_updateGradInput(backend.library_state,
@@ -113,7 +147,11 @@ def backward(ctx, grad_output, _indices_grad=None):
  ctx.padding[1], ctx.padding[0],
  ctx.dilation[1], ctx.dilation[0],
  ctx.ceil_mode)
- return grad_input, None, None, None, None, None, None
+ return grad_input
+
+ @staticmethod
+ def backward(ctx, ggI, _ggIndices=None):
+ raise ValueError("MaxPool2d cannot be differentiated twice")
 
 
 class MaxPool3d(Function):
@@ -146,13 +184,28 @@ def forward(ctx, input, kernel_size, stride=None, padding=0, dilation=1,
  return output
 
  @staticmethod
- @once_differentiable
  def backward(ctx, grad_output, _indices_grad=None):
  if ctx.return_indices:
- input, indices = ctx.saved_tensors
+ input, indices = ctx.saved_variables
  else:
- input, = ctx.saved_tensors
+ input, = ctx.saved_variables
  indices = ctx.indices
+ grad_input = MaxPool3dBackward.apply(input, indices, grad_output, ctx.kernel_size, ctx.stride,
+ ctx.padding, ctx.dilation, ctx.return_indices, ctx.ceil_mode)
+ return grad_input, None, None, None, None, None, None
+
+
+class MaxPool3dBackward(Function):
+
+ @staticmethod
+ def forward(ctx, input, indices, grad_output, kernel_size, stride, padding, dilation,
+ return_indices, ceil_mode):
+ ctx.kernel_size = kernel_size
+ ctx.stride = stride
+ ctx.padding = padding
+ ctx.dilation = dilation
+ ctx.return_indices = return_indices
+ ctx.ceil_mode = ceil_mode
  grad_input = grad_output.new()
  backend = type2backend[type(input)]
  backend.VolumetricDilatedMaxPooling_updateGradInput(backend.library_state,
@@ -163,7 +216,11 @@ def backward(ctx, grad_output, _indices_grad=None):
  ctx.padding[0], ctx.padding[2], ctx.padding[1],
  ctx.dilation[0], ctx.dilation[2], ctx.dilation[1],
  ctx.ceil_mode)
- return grad_input, None, None, None, None, None, None
+ return grad_input
+
+ @staticmethod
+ def backward(ctx, ggI, _ggIndices=None):
+ raise ValueError("MaxPool3d cannot be differentiated twice")
 
 
 class MaxUnpool2d(Function):
@@ -382,11 +439,11 @@ def forward(ctx, input, kernel_size, stride=None):
  @staticmethod
  def backward(ctx, grad_output):
  input, = ctx.saved_variables
- grad_input = AvgPool3dBackwards.apply(input, grad_output, ctx.kernel_size, ctx.stride)
+ grad_input = AvgPool3dBackward.apply(input, grad_output, ctx.kernel_size, ctx.stride)
  return grad_input, None, None
 
 
-class AvgPool3dBackwards(Function):
+class AvgPool3dBackward(Function):
 
  @staticmethod
  def forward(ctx, input, grad_output, kernel_size, stride):
@@ -408,6 +465,7 @@ def backward(ctx, ggI):
  ggO = AvgPool3d.apply(ggI, ctx.kernel_size, ctx.stride)
  return gI, ggO, None, None
 
+
 class AdaptiveMaxPool1d(Function):
 
  @staticmethod
@@ -553,9 +611,13 @@ def backward(ctx, grad_output):
 _all_functions.append(AvgPool2d)
 _all_functions.append(AvgPool2dBackward)
 _all_functions.append(AvgPool3d)
+_all_functions.append(AvgPool3dBackward)
 _all_functions.append(MaxPool1d)
+_all_functions.append(MaxPool1dBackward)
 _all_functions.append(MaxPool2d)
+_all_functions.append(MaxPool2dBackward)
 _all_functions.append(MaxPool3d)
+_all_functions.append(MaxPool3dBackward)
 _all_functions.append(MaxUnpool2d)
 _all_functions.append(MaxUnpool3d)
 _all_functions.append(FractionalMaxPool2d)