-
Notifications
You must be signed in to change notification settings - Fork 159
Expand file tree
/
Copy pathPointwise.lua
More file actions
94 lines (82 loc) · 2.76 KB
/
Copy pathPointwise.lua
File metadata and controls
94 lines (82 loc) · 2.76 KB
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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
local Pointwise, parent = torch.class('cudnn._Pointwise','nn.Module')
local errcheck = cudnn.errcheck
local ffi = require 'ffi'
function Pointwise:__init(inplace)
parent.__init(self)
self.inplace = inplace or false
end
function Pointwise:createIODescriptors(input)
assert(self.mode, 'mode is not set. (trying to use base class?)');
assert(input:isContiguous(), 'Non-contiguous inputs not supported yet');
if not self.inplace then
self.output:resizeAs(input)
end
if not self.activDesc then
self.activDesc = ffi.new('struct cudnnActivationStruct*[1]')
errcheck('cudnnCreateActivationDescriptor', self.activDesc)
errcheck('cudnnSetActivationDescriptor', self.activDesc[0], self.mode, 'CUDNN_PROPAGATE_NAN', self.ceiling or 0.0);
local function destroyADesc(a)
if (a[0]) then
errcheck('cudnnDestroyActivationDescriptor', a[0]);
a[0] = nil
end
end
ffi.gc(self.activDesc, destroyADesc)
end
local nElem = input:nElement()
self.nElem = self.nElem or nElem -- this goes to the second branch only once
if self.iDesc and nElem == self.nElem then return end
self.nElem = nElem
self.iDesc = cudnn.toDescriptor(input:view(1,1,1,nElem))
end
function Pointwise:updateOutput(input)
self:createIODescriptors(input)
if self.inplace then self.output:set(input) end
errcheck('cudnnActivationForward',
cudnn.getHandle(), self.activDesc[0],
cudnn.scalar(input, 1),
self.iDesc[0], input:data(),
cudnn.scalar(input, 0),
self.iDesc[0], self.output:data());
return self.output
end
function Pointwise:updateGradInput(input, gradOutput)
if not gradOutput:isContiguous() then
self._gradOutput = self._gradOutput or gradOutput.new()
self._gradOutput:resizeAs(gradOutput):copy(gradOutput)
gradOutput = self._gradOutput
end
self:createIODescriptors(input)
if self.inplace then
self.output:set(input);
self.gradInput:set(gradOutput)
else
self.gradInput:resizeAs(input)
end
errcheck('cudnnActivationBackward',
cudnn.getHandle(), self.activDesc[0],
cudnn.scalar(input, 1),
self.iDesc[0], self.output:data(),
self.iDesc[0], gradOutput:data(),
self.iDesc[0], input:data(),
cudnn.scalar(input, 0),
self.iDesc[0], self.gradInput:data());
return self.gradInput
end
function Pointwise:clearDesc()
self.iDesc = nil
self.activDesc = nil
end
function Pointwise:write(f)
self:clearDesc()
local var = {}
for k,v in pairs(self) do
var[k] = v
end
f:writeObject(var)
end
function Pointwise:clearState()
self:clearDesc()
nn.utils.clear(self, '_gradOutput')
return parent.clearState(self)
end