| import math |
| import torch |
| from .Module import Module |
| from .Sequential import Sequential |
| from .SpatialZeroPadding import SpatialZeroPadding |
| from .SpatialConvolution import SpatialConvolution |
| from .SpatialConvolutionMap import SpatialConvolutionMap |
| from .Replicate import Replicate |
| from .CSubTable import CSubTable |
| from .CDivTable import CDivTable |
| from .utils import clear |
| import warnings |
| |
| |
| class SpatialSubtractiveNormalization(Module): |
| |
| def __init__(self, nInputPlane=1, kernel=None): |
| super(SpatialSubtractiveNormalization, self).__init__() |
| |
| # get args |
| self.nInputPlane = nInputPlane |
| if kernel is None: |
| kernel = torch.Tensor(9, 9).fill_(1) |
| self.kernel = kernel |
| kdim = self.kernel.ndimension() |
| |
| # check args |
| if kdim != 2 and kdim != 1: |
| raise ValueError('SpatialSubtractiveNormalization averaging kernel must be 2D or 1D') |
| |
| if (self.kernel.size(0) % 2) == 0 or (kdim == 2 and (self.kernel.size(1) % 2) == 0): |
| raise ValueError('SpatialSubtractiveNormalization averaging kernel must have ODD dimensions') |
| |
| # normalize kernel |
| self.kernel.div_(self.kernel.sum() * self.nInputPlane) |
| |
| # padding values |
| padH = int(math.floor(self.kernel.size(0) / 2)) |
| padW = padH |
| if kdim == 2: |
| padW = int(math.floor(self.kernel.size(1) / 2)) |
| |
| # create convolutional mean extractor |
| self.meanestimator = Sequential() |
| self.meanestimator.add(SpatialZeroPadding(padW, padW, padH, padH)) |
| if kdim == 2: |
| self.meanestimator.add(SpatialConvolution(self.nInputPlane, 1, self.kernel.size(1), self.kernel.size(0))) |
| else: |
| # TODO: map |
| self.meanestimator.add(SpatialConvolutionMap( |
| SpatialConvolutionMap.maps.oneToOne(self.nInputPlane), self.kernel.size(0), 1)) |
| self.meanestimator.add(SpatialConvolution(self.nInputPlane, 1, 1, self.kernel.size(0))) |
| |
| self.meanestimator.add(Replicate(self.nInputPlane, 0)) |
| |
| # set kernel and bias |
| if kdim == 2: |
| for i in range(self.nInputPlane): |
| self.meanestimator.modules[1].weight[0][i] = self.kernel |
| self.meanestimator.modules[1].bias.zero_() |
| else: |
| for i in range(self.nInputPlane): |
| self.meanestimator.modules[1].weight[i].copy_(self.kernel, broadcast=False) |
| self.meanestimator.modules[2].weight[0][i].copy_(self.kernel, broadcast=False) |
| |
| self.meanestimator.modules[1].bias.zero_() |
| self.meanestimator.modules[2].bias.zero_() |
| |
| # other operation |
| self.subtractor = CSubTable() |
| self.divider = CDivTable() |
| |
| # coefficient array, to adjust side effects |
| self.coef = torch.Tensor(1, 1, 1) |
| |
| self.ones = None |
| self._coef = None |
| |
| def updateOutput(self, input): |
| # compute side coefficients |
| dim = input.dim() |
| if (input.dim() + 1 != self.coef.dim() or |
| (input.size(dim - 1) != self.coef.size(dim - 1)) or |
| (input.size(dim - 2) != self.coef.size(dim - 2))): |
| if self.ones is None: |
| self.ones = input.new() |
| if self._coef is None: |
| self._coef = self.coef.new() |
| |
| self.ones.resize_as_(input[0:1]).fill_(1) |
| coef = self.meanestimator.updateOutput(self.ones).squeeze(0) |
| self._coef.resize_as_(coef).copy_(coef) # make contiguous for view |
| size = list(coef.size()) |
| size = [input.size(0)] + size |
| self.coef = self._coef.view(1, *self._coef.size()).expand(*size) |
| |
| # compute mean |
| self.localsums = self.meanestimator.updateOutput(input) |
| self.adjustedsums = (self.divider.updateOutput( |
| [self.localsums, self.coef.contiguous().view_as(self.localsums)])) |
| self.output = self.subtractor.updateOutput([input, self.adjustedsums.contiguous().view_as(input)]) |
| |
| return self.output |
| |
| def updateGradInput(self, input, gradOutput): |
| # resize grad |
| self.gradInput.resize_as_(input).zero_() |
| |
| # backprop through all modules |
| gradsub = self.subtractor.updateGradInput([input, self.adjustedsums.contiguous().view_as(input)], gradOutput) |
| graddiv = (self.divider.updateGradInput( |
| [self.localsums, self.coef.contiguous().view_as(self.localsums)], gradsub[1])) |
| size = self.meanestimator.updateGradInput(input, graddiv[0]).size() |
| self.gradInput.add_(self.meanestimator.updateGradInput(input, graddiv[0])) |
| self.gradInput.add_(gradsub[0]) |
| |
| return self.gradInput |
| |
| def clearState(self): |
| clear(self, 'ones', '_coef') |
| self.meanestimator.clearState() |
| return super(SpatialSubtractiveNormalization, self).clearState() |
