Make it easier to check with cuda.
[dagnn.git] / dagnn.lua
1
2 --[[
3
4    Copyright (c) 2016 Idiap Research Institute, http://www.idiap.ch/
5    Written by Francois Fleuret <francois.fleuret@idiap.ch>
6
7    This file is free software: you can redistribute it and/or modify
8    it under the terms of the GNU General Public License version 3 as
9    published by the Free Software Foundation.
10
11    It is distributed in the hope that it will be useful, but WITHOUT
12    ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
13    or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
14    License for more details.
15
16    You should have received a copy of the GNU General Public License
17    along with this file.  If not, see <http://www.gnu.org/licenses/>.
18
19 ]]--
20
21 require 'torch'
22 require 'nn'
23
24 local DAG, parent = torch.class('nn.DAG', 'nn.Container')
25
26 function DAG:__init()
27    parent.__init(self)
28    -- Nodes are indexed by the module they contain
29    self.node = {}
30 end
31
32 -- Apply f on t recursively; use the corresponding elements from args
33 -- (i.e. same keys) as second parameter to f when available; return
34 -- the results from f, organized in a similarly nested table.
35 function DAG:nestedApply(f, t, args)
36    if torch.type(t) == 'table' then
37       local result = {}
38       for k, s in pairs(t) do
39          result[k] = self:nestedApply(f, s, args and args[k])
40       end
41       return result
42    else
43       return f(t, args)
44    end
45 end
46
47 function DAG:createNode(nnm)
48    if not self.node[nnm] then
49       self:add(nnm) -- Add it to the object as a Container
50       local node = {}
51       node.succ = {}
52       node.pred = {}
53       node.index = #self.modules
54       self.node[nnm] = node
55    end
56 end
57
58 function DAG:putInOrder()
59    if self.sorted then
60       return
61    end
62
63    local distance = {}
64    self:nestedApply(
65       function(m) distance[m] = 1 end,
66       self.inputModules
67    )
68
69    local nc
70    local nl = 0
71    repeat
72       nc = 0
73       for nnma, node in pairs(self.node) do
74          for _, nnmb in pairs(node.succ) do
75             if distance[nnma] and (not distance[nnmb] or distance[nnmb] < distance[nnma] + 1) then
76                distance[nnmb] = distance[nnma] + 1
77                nc = nc + 1
78             end
79          end
80       end
81       assert(nl < #self.modules, 'Cycle detected in the graph.')
82       nl = nl + 1
83    until nc == 0
84
85    for _, nnm in pairs(self.modules) do
86       assert(distance[nnm], 'Some modules are not connected to inputs')
87    end
88
89    self.sorted = {}
90    for m, d in pairs(distance) do
91       table.insert(self.sorted, { distance = d, nnm = m })
92    end
93
94    table.sort(self.sorted, function(a, b) return a.distance < b.distance end)
95
96    for i, a in ipairs(self.sorted) do self.sorted[i] = a.nnm end
97 end
98
99 -- This accumulates x in a, where they are both nested tables of
100 -- tensors with same structures / keys. If first is true, set a = x
101 -- (in which case a can be nil) otherwise a = a + x. The behavior is
102 -- undefined if a and x do not have the exact same structure.
103 function DAG:nestedAccTensor(a, x, first)
104    if torch.type(x) == 'table' then
105       local b = {}
106       for i in pairs(x) do
107          b[i] = self:nestedAccTensor(a[i], x[i], first)
108       end
109       a = b
110    else
111       if first then
112          if a then
113             a:resizeAs(x):copy(x)
114          else
115             a = x:clone()
116          end
117       else
118          a:add(x)
119       end
120    end
121    return a
122 end
123
124 function DAG:updateGradOutput(node)
125    local gradInputSucc = node.gradInputSucc
126    if #gradInputSucc == 1 then
127       node.gradOutput = gradInputSucc[1]
128    elseif #gradInputSucc > 1 then
129       for k = 1, #gradInputSucc do
130          node.gradOutput = self:nestedAccTensor(node.gradOutput, gradInputSucc[k], k == 1)
131       end
132    end
133 end
134
135 ----------------------------------------------------------------------
136
137 -- Connect a sequence of modules
138 function DAG:connect(...)
139    self.sorted = nil
140    local prev
141    for _, nnm in pairs({...}) do
142       self:createNode(nnm)
143       if prev then
144          table.insert(self.node[nnm].pred, prev)
145          table.insert(self.node[prev].succ, nnm)
146       end
147       prev = nnm
148    end
149 end
150
151 function DAG:setInput(i)
152    self.sorted = nil
153    self.inputModules = i
154    self:nestedApply(
155       function(nnm)
156          assert(#self.node[nnm].succ > 0, 'Input modules must have outgoing edges.')
157          assert(#self.node[nnm].pred == 0, 'Input modules cannot have incoming edges.')
158       end,
159       self.inputModules
160    )
161 end
162
163 function DAG:setOutput(o)
164    self.sorted = nil
165    self.outputModules = o
166    self:nestedApply(
167       function(nnm)
168          assert(#self.node[nnm].pred > 0, 'Output module must have incoming edges.')
169          assert(#self.node[nnm].succ == 0, 'Output module cannot have outgoing edges.')
170       end,
171       self.outputModules
172    )
173 end
174
175 function DAG:print()
176    self:putInOrder()
177
178    for i, d in ipairs(self.sorted) do
179       print('#' .. i .. ' -> ' .. torch.type(d))
180    end
181 end
182
183 ----------------------------------------------------------------------
184
185 function DAG:saveDot(filename)
186    local file = (filename and io.open(filename, 'w')) or io.stdout
187
188    local function writeNestedCluster(prefix, list, indent)
189       local indent = indent or ''
190       if torch.type(list) == 'table' then
191          file:write(indent .. '  subgraph cluster_' .. prefix .. ' {\n');
192          for k, x in pairs(list) do
193             writeNestedCluster(prefix .. '_' .. k, x, '  ' .. indent)
194          end
195          file:write(indent .. '  }\n');
196       else
197          file:write(indent .. '  ' .. self.node[list].index .. ' [color=red]\n')
198       end
199    end
200
201    file:write('digraph {\n')
202
203    file:write('\n')
204
205    writeNestedCluster('input', self.inputModules)
206    writeNestedCluster('output', self.outputModules)
207
208    file:write('\n')
209
210    for nnmb, node in pairs(self.node) do
211       file:write(
212          '  '
213             .. node.index
214             .. ' [shape=box,label=\"' .. torch.type(nnmb) .. '\"]'
215             .. '\n'
216       )
217
218       for i, nnma in pairs(node.pred) do
219          local decoration = ''
220          if #node.pred > 1 then
221             -- decoration = ' [headlabel=\"' .. i .. '\"]'
222             decoration = ' [label=\"' .. i .. '\"]'
223          end
224          file:write(
225             '  '
226                .. self.node[nnma].index
227                .. ' -> '
228                .. self.node[nnmb].index
229                .. decoration
230                .. '\n'
231          )
232       end
233
234       file:write('\n')
235    end
236
237    file:write('}\n')
238
239 end
240
241 ----------------------------------------------------------------------
242
243 function DAG:updateOutput(input)
244    self:putInOrder()
245
246    self:nestedApply(
247       function(nnm, i)
248          local node = self.node[nnm]
249          node.input = i
250          self:rethrowErrors(nnm, node.index, 'updateOutput', i)
251       end,
252       self.inputModules,
253       input
254    )
255
256    for _, nnm in ipairs(self.sorted) do
257       local node = self.node[nnm]
258       local pred = node.pred
259       if #pred > 0 then
260          local i
261          if #pred == 1 then
262             i = pred[1].output
263          elseif #pred > 1 then
264             i = {}
265             for k = 1, #pred do
266                i[k] = pred[k].output
267             end
268          end
269          node.input = i
270          self:rethrowErrors(nnm, node.index, 'updateOutput', i)
271       end
272    end
273
274    self.output = self:nestedApply(
275       function(m) return m.output end,
276       self.outputModules
277    )
278
279    return self.output
280 end
281
282 function DAG:updateGradInput(input, gradOutput)
283    assert(self.sorted, 'There has been a structure change before a DAG:updateGradInput')
284
285    self:nestedApply(
286       function(nnm, go)
287          local node = self.node[nnm]
288          node.gradOutput = go
289          self:rethrowErrors(nnm, node.index, 'updateGradInput', node.input, go)
290       end,
291       self.outputModules, gradOutput
292    )
293
294    self:nestedApply(
295       function(nnm, i) self.node[nnm].input = i end,
296       self.inputModules, input
297    )
298
299    for _, node in pairs(self.node) do
300       node.gradInputSucc = {}
301    end
302
303    for k = #self.sorted, 1, -1 do
304       local nnm = self.sorted[k]
305       local node = self.node[nnm]
306       local pred = node.pred
307
308       if #node.gradInputSucc > 0 then
309          self:updateGradOutput(node)
310          self:rethrowErrors(nnm, node.index, 'updateGradInput', node.input, node.gradOutput)
311       end
312
313       -- We fill the gradInputSucc of our predecessors
314       if #pred == 1 then
315          table.insert(self.node[pred[1]].gradInputSucc, nnm.gradInput)
316       elseif #pred > 1 then
317          assert(torch.type(nnm.gradInput) == 'table',
318                 'Should have a table gradInput since it has multiple predecessors')
319          for n = 1, #pred do
320             table.insert(self.node[pred[n]].gradInputSucc, nnm.gradInput[n])
321          end
322       end
323    end
324
325    self.gradInput = self:nestedApply(
326       function(m) return m.gradInput end,
327       self.inputModules
328    )
329
330    return self.gradInput
331 end
332
333 function DAG:accGradParameters(input, gradOutput, scale)
334    assert(self.sorted, 'There has been a structure change before a DAG:accGradParameters')
335
336    self:nestedApply(
337       function(nnm, go) self.node[nnm].gradOutput = go end,
338       self.outputModules, gradOutput
339    )
340
341    self:nestedApply(
342       function(nnm, i) self.node[nnm].input = i end,
343       self.inputModules, input
344    )
345
346    for k = 1, #self.modules do
347       local nnm = self.modules[k]
348       local node = self.node[nnm]
349       self:rethrowErrors(nnm, k, 'accGradParameters', node.input, node.gradOutput, scale)
350    end
351 end
352
353 function DAG:clearState()
354    self.sorted = nil
355    for _, node in pairs(self.node) do
356       node.input = nil
357       node.gradInputSucc = nil
358       node.gradOutput = nil
359    end
360    return parent.clearState(self)
361 end