-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathSplitTableMultiple.lua
46 lines (39 loc) · 1.47 KB
/
SplitTableMultiple.lua
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
36
37
38
39
40
41
42
43
44
45
local SplitTableMultiple, parent = torch.class('nn.SplitTableMultiple', 'nn.Module')
function SplitTableMultiple:__init(dimension, nInputDims, blockSize)
parent.__init(self)
self.dimension = dimension
self.nInputDims = nInputDims
self.blockSize = blockSize or 1
end
function SplitTableMultiple:_getPositiveDimension(input)
local dimension = self.dimension
if dimension < 0 then
dimension = input:dim() + dimension + 1
elseif self.nInputDims and input:dim()==(self.nInputDims+1) then
dimension = dimension + 1
end
return dimension
end
function SplitTableMultiple:updateOutput(input)
local dimension = self:_getPositiveDimension(input)
local slices = input:size(dimension)
assert(math.fmod(slices,self.blockSize)==0,'Dimension must be multiple of blockSize')
local currentOutput= {}
for i=1,slices,self.blockSize do
currentOutput[#currentOutput+1] = input:narrow(dimension,i,self.blockSize):contiguous()
end
self.output = currentOutput
return self.output
end
function SplitTableMultiple:updateGradInput(input, gradOutput)
local dimension = self:_getPositiveDimension(input)
local slices = #gradOutput -- input:size(dimension)
if self.gradInput then
self.gradInput:resizeAs(input)
for i=1,slices do
local currentGradInput = gradOutput[i];
self.gradInput:narrow(dimension,(i-1)*self.blockSize+1,self.blockSize):copy(currentGradInput)
end
end
return self.gradInput
end