--- a +++ b/util/ReverseSequence.lua @@ -0,0 +1,55 @@ +------------------------------------------------------------------------ +-- Adapted from github.com/jcjohnson/torch-rnn/pull/66/commits/5e30c1d54dc9ed1d152e4a55e9c438775f623ef7 + +--[[ ReverseSequence ]] -- +-- Reverses a sequence on a given dimension. +-- Example: Given a tensor of torch.Tensor({{1,2,3,4,5}, {6,7,8,9,10}) +-- nn.ReverseSequence(1):forward(tensor) would give: torch.Tensor({{6,7,8,9,10},{1,2,3,4,5}}) +------------------------------------------------------------------------ +local ReverseSequence, parent = torch.class("nn.ReverseSequence", "nn.Module") + +function ReverseSequence:__init(dim,gpu) + parent.__init(self) + self.output = torch.Tensor() + self.gradInput = torch.Tensor() + self.outputIndices = torch.LongTensor() + self.gradIndices = torch.LongTensor() + self.typ = 'torch.CudaTensor' + if gpu and (gpu < 1) then + self.typ = 'torch.LongTensor' + end +end + +function ReverseSequence:reverseOutput(input) + self.output:resizeAs(input) + self.outputIndices:resize(input:size()) + local T = input:size(1) + for x = 1, T do + self.outputIndices:narrow(1, x, 1):fill(T - x + 1) + end + self.output:gather(input, 1, self.outputIndices:type(self.typ)) +end + +function ReverseSequence:updateOutput(input) + input = input:transpose(1, 2) + self:reverseOutput(input) + self.output = self.output:transpose(1, 2) + return self.output +end + +function ReverseSequence:reverseGradOutput(gradOutput) + self.gradInput:resizeAs(gradOutput) + self.gradIndices:resize(gradOutput:size()) + local T = gradOutput:size(1) + for x = 1, T do + self.gradIndices:narrow(1, x, 1):fill(T - x + 1) + end + self.gradInput:gather(gradOutput, 1, self.gradIndices:type(self.typ)) +end + +function ReverseSequence:updateGradInput(inputTable, gradOutput) + gradOutput = gradOutput:transpose(1, 2) + self:reverseGradOutput(gradOutput) + self.gradInput = self.gradInput:transpose(1, 2) + return self.gradInput +end