| import math |
| |
| import torch |
| from .Module import Module |
| from .utils import clear, contiguousView |
| |
| |
| class CMul(Module): |
| |
| def __init__(self, *args): |
| super(CMul, self).__init__() |
| |
| if len(args) == 1 and isinstance(args[0], torch.Size): |
| self.size = args[0] |
| else: |
| self.size = torch.Size(args) |
| |
| self.weight = torch.Tensor(self.size) |
| self.gradWeight = torch.Tensor(self.size) |
| self.output.resize_(self.size) |
| self.reset() |
| |
| self._output = None |
| self._weight = None |
| self._expand = None |
| self._repeat = None |
| self._gradOutput = None |
| self._gradInput = None |
| self._input = None |
| self._gradWeight = None |
| self._sum = None |
| |
| def reset(self, stdv=None): |
| if stdv is not None: |
| stdv = stdv * math.sqrt(3) |
| else: |
| stdv = 1. / math.sqrt(self.weight.nelement()) |
| |
| self.weight.uniform_(-stdv, stdv) |
| |
| def updateOutput(self, input): |
| # lazy-initialize |
| if self._output is None: |
| self._output = input.new() |
| self._weight = input.new() |
| self._expand = input.new() |
| self._repeat = input.new() |
| |
| self.output.resize_as_(input).copy_(input) |
| batchSize = input.size(0) |
| # TODO: expand_as_, view_ |
| self._output = self.output.view(batchSize, -1) |
| self._weight = self.weight.view(1, -1) |
| self._expand = self._weight.expand_as(self._output) |
| |
| if torch.typename(input) == 'torch.cuda.FloatTensor': |
| self._repeat.resize_as_(self._expand).copy_(self._expand) |
| self._output.mul_(self._repeat) |
| else: |
| self._output.mul_(self._expand) |
| |
| return self.output |
| |
| def updateGradInput(self, input, gradOutput): |
| if self.gradInput is None: |
| return |
| |
| if self._gradOutput is None: |
| self._gradOutput = input.new() |
| self._gradInput = input.new() |
| |
| self.gradInput.resize_as_(input).zero_() |
| batchSize = input.size(0) |
| contiguousView(self._gradOutput, gradOutput, batchSize, -1) |
| contiguousView(self._gradInput, self.gradInput, batchSize, -1) |
| self._weight = self.weight.view(1, -1) |
| self._expand = self._weight.expand_as(self._gradOutput) |
| |
| if torch.typename(input) == 'torch.cuda.FloatTensor': |
| self._repeat.resize_as_(self._expand).copy_(self._expand) |
| self._gradInput.addcmul_(1, self._repeat, self._gradOutput) |
| else: |
| self._gradInput.addcmul_(1, self._expand, self._gradOutput) |
| |
| return self.gradInput |
| |
| def accGradParameters(self, input, gradOutput, scale=1): |
| if self._input is None: |
| self._input = input.new() |
| self._gradWeight = input.new() |
| self._sum = input.new() |
| |
| batchSize = input.size(0) |
| contiguousView(self._input, input, batchSize, -1) |
| contiguousView(self._gradOutput, gradOutput, batchSize, -1) |
| self._gradWeight = self.gradWeight.view(1, -1) |
| |
| torch.mul(self._input, self._gradOutput, out=self._repeat) |
| torch.sum(self._repeat, 0, True, out=self._sum) |
| self._gradWeight.add_(scale, self._sum) |
| |
| def type(self, type=None, tensorCache=None): |
| if type: |
| self.clearState() |
| return super(CMul, self).type(type, tensorCache) |
| |
| def clearState(self): |
| clear(self, [ |
| '_input', |
| '_output', |
| '_weight', |
| '_gradWeight', |
| '_expand', |
| '_repeat', |
| '_sum', |
| ]) |
| return super(CMul, self).clearState() |
