blob: cceb2b643ea68e80bb96e9488f689910874fc1c2 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
|
local Transpose, parent = torch.class('nn.Transpose', 'nn.Module')
-- transpose dimensions:
-- n = nn.Transpose({1,4},{1,3})
-- will transpose dims 1 and 4, then 1 and 3...
function Transpose:__init(...)
parent.__init(self)
self.permutations = {...}
self.numInputDims = nil
end
function Transpose:setNumInputDims(numInputDims)
self.numInputDims = numInputDims
return self
end
function Transpose:updateOutput(input)
local offset = self.numInputDims and input:nDimension()-self.numInputDims or 0
for _,perm in ipairs(self.permutations) do
input = input:transpose(perm[1]+offset,perm[2]+offset)
end
self.output:resizeAs(input):copy(input)
return self.output
end
function Transpose:updateGradInput(input, gradOutput)
for i = #self.permutations,1,-1 do
local perm = self.permutations[i]
local offset = self.numInputDims and input:nDimension()-self.numInputDims or 0
gradOutput = gradOutput:transpose(perm[1]+offset,perm[2]+offset)
end
self.gradInput:resizeAs(gradOutput):copy(gradOutput)
return self.gradInput
end
|
