blob: a6ebaa3213d17b94c46cd2628795387199fd4b59 (
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
36
37
38
39
40
41
42
43
44
45
|
local Narrow, parent = torch.class('nn.Narrow', 'nn.Module')
function Narrow:__init(dimension,offset,length)
parent.__init(self)
self.dimension=dimension
self.index=offset
self.length=length or 1
if not dimension or not offset then
error('nn.Narrow(dimension, offset, length)')
end
end
function Narrow:updateOutput(input)
local dim = self.dimension < 0 and input:dim() + self.dimension + 1 or self.dimension
local length = self.length
if length < 0 then
length = input:size(dim) - self.index + self.length + 2
end
local index = self.index
if self.index < 0 then
index = 1
length = input:size(dim) - length
end
local output=input:narrow(dim, index, length)
self.output = self.output:typeAs(output)
self.output:resizeAs(output):copy(output)
return self.output
end
function Narrow:updateGradInput(input, gradOutput)
local dim = self.dimension < 0 and input:dim() + self.dimension + 1 or self.dimension
local length = self.length
if length < 0 then
length = input:size(dim) - self.index + self.length + 2
end
local index = self.index
if self.index < 0 then
index = 1
length = input:size(dim) - length
end
self.gradInput = self.gradInput:typeAs(input)
self.gradInput:resizeAs(input):zero()
self.gradInput:narrow(dim,index,length):copy(gradOutput)
return self.gradInput
end
|
