| import torch |
| from .Module import Module |
| from .utils import clear |
| |
| |
| class CosineDistance(Module): |
| |
| def __init__(self, ): |
| super(CosineDistance, self).__init__() |
| self.gradInput = [torch.Tensor(), torch.Tensor()] |
| |
| self._input1 = None |
| self._input2 = None |
| self.buffer = None |
| self.w1 = None |
| self.w22 = None |
| self.w = None |
| self.w32 = None |
| self.ones = None |
| |
| def _makeContiguous(self, input1, input2): |
| if not input1.is_contiguous(): |
| if self._input1 is None: |
| self._input1 = input1.new() |
| self._input1.resize_as_(input1).copy_(input1) |
| input1 = self._input1 |
| |
| if not input2.is_contiguous(): |
| if self._input2 is None: |
| self._input2 = input2.new() |
| self._input2.resize_as_(input2).copy_(input2) |
| input2 = self._input2 |
| |
| return input1, input2 |
| |
| def updateOutput(self, input): |
| input1, input2 = input[0], input[1] |
| input1, input2 = self._makeContiguous(input1, input2) |
| |
| if self.buffer is None: |
| self.buffer = input1.new() |
| self.w1 = input1.new() |
| self.w22 = input1.new() |
| self.w = input1.new() |
| self.w32 = input1.new() |
| self.ones = input1.new() |
| |
| torch.mul(input1, input2, out=self.buffer) |
| torch.sum(self.buffer, 1, out=self.w1, keepdim=True) |
| |
| epsilon = 1e-12 |
| torch.mul(input1, input1, out=self.buffer) |
| torch.sum(self.buffer, 1, out=self.w22, keepdim=True).add_(epsilon) |
| self.w22.reciprocal_() |
| self.w.resize_as_(self.w22).copy_(self.w22) |
| |
| torch.mul(input2, input2, out=self.buffer) |
| torch.sum(self.buffer, 1, out=self.w32, keepdim=True).add_(epsilon) |
| self.w32.reciprocal_() |
| self.w.mul_(self.w32) |
| self.w.sqrt_() |
| |
| torch.mul(self.w1, self.w, out=self.output) |
| self.output.resize_(input1.size(0)) |
| |
| return self.output |
| |
| def updateGradInput(self, input, gradOutput): |
| v1 = input[0] |
| v2 = input[1] |
| v1, v2 = self._makeContiguous(v1, v2) |
| |
| if len(self.gradInput) != 2: |
| if self.gradInput[0] is None: |
| self.gradInput[0] = v1.new() |
| if self.gradInput[1] is None: |
| self.gradInput[1] = v1.new() |
| self.gradInput = self.gradInput[:2] |
| |
| gw1 = self.gradInput[0] |
| gw2 = self.gradInput[1] |
| gw1.resize_as_(v1).copy_(v2) |
| gw2.resize_as_(v1).copy_(v1) |
| |
| torch.mul(self.w1, self.w22, out=self.buffer) |
| gw1.addcmul_(-1, self.buffer.expand_as(v1), v1) |
| gw1.mul_(self.w.expand_as(v1)) |
| |
| torch.mul(self.w1, self.w32, out=self.buffer) |
| gw2.addcmul_(-1, self.buffer.expand_as(v1), v2) |
| gw2.mul_(self.w.expand_as(v1)) |
| |
| go = gradOutput.contiguous().view(-1, 1).expand_as(v1) |
| gw1.mul_(go) |
| gw2.mul_(go) |
| |
| return self.gradInput |
| |
| def clearState(self): |
| clear(self, [ |
| 'buffer', |
| 'w1', |
| 'w22', |
| 'w', |
| 'w32', |
| 'ones', |
| ]) |
| return super(CosineDistance, self).clearState() |
