| import torch |
| from .Module import Module |
| from .utils import clear |
| |
| |
| class Sum(Module): |
| |
| def __init__(self, dimension=0, sizeAverage=False): |
| super(Sum, self).__init__() |
| self.dimension = dimension |
| self.sizeAverage = sizeAverage |
| self._gradOutput = None |
| |
| def _getPositiveDimension(self, input): |
| dimension = self.dimension |
| if dimension < 0: |
| dimension = input.dim() + dimension |
| return dimension |
| |
| def updateOutput(self, input): |
| dimension = self._getPositiveDimension(input) |
| |
| torch.sum(input, dimension, out=self.output, keepdim=True) |
| if self.sizeAverage: |
| self.output.div_(input.size(dimension)) |
| if self.output.dim() > 1: |
| self.output.set_(self.output.select(dimension, 0)) |
| |
| return self.output |
| |
| def updateGradInput(self, input, gradOutput): |
| dimension = self._getPositiveDimension(input) |
| # zero-strides don't work with MKL/BLAS, so |
| # don't set self.gradInput to zero-stride tensor. |
| # Instead, do a deepcopy. |
| size = list(input.size()) |
| size[dimension] = 1 |
| if not gradOutput.is_contiguous(): |
| if self._gradOutput is None: |
| self._gradOutput = gradOutput.new() |
| self._gradOutput.resize_as_(gradOutput).copy_(gradOutput) |
| gradOutput = self._gradOutput |
| |
| gradOutput = gradOutput.view(*size) |
| self.gradInput.resize_as_(input) |
| self.gradInput.copy_(gradOutput.expand_as(input)) |
| if self.sizeAverage: |
| self.gradInput.div_(input.size(dimension)) |
| |
| return self.gradInput |
| |
| def clearState(self): |
| clear(self, '_gradOutput') |
| return super(Sum, self).clearState() |
