cf452338f3749dded66500717117464ee6c82e70
[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(function(m) distance[m] = 1 end, self.inputModules)
65
66    local nc
67    repeat
68       nc = 0
69       for nnma, node in pairs(self.node) do
70          for _, nnmb in pairs(node.succ) do
71             if distance[nnma] and (not distance[nnmb] or distance[nnmb] < distance[nnma] + 1) then
72                distance[nnmb] = distance[nnma] + 1
73                nc = nc + 1
74             end
75          end
76       end
77    until nc == 0
78
79    for _, nnm in pairs(self.modules) do
80       assert(distance[nnm], 'Some modules are not connected to inputs')
81    end
82
83    self.sorted = {}
84    for m, d in pairs(distance) do
85       table.insert(self.sorted, { distance = d, nnm = m })
86    end
87
88    table.sort(self.sorted, function(a, b) return a.distance < b.distance end)
89
90    for i, a in ipairs(self.sorted) do self.sorted[i] = a.nnm end
91 end
92
93 -- This accumulates x in a where they are both nested tables of
94 -- tensors. If first is true, set a = x. Behavior is undefined if a
95 -- and x do not have the exact same structure.
96 function DAG:nestedAccTensor(a, x, first)
97    if torch.type(x) == 'table' then
98       local b = {}
99       for i in pairs(x) do
100          b[i] = self:nestedAccTensor(a[i], x[i], first)
101       end
102       a = b
103    else
104       if first then
105          if a then
106             a:resizeAs(x):copy(x)
107          else
108             a = x:clone()
109          end
110       else
111          a:add(x)
112       end
113    end
114    return a
115 end
116
117 function DAG:updateGradOutput(node)
118    local gradInputSucc = node.gradInputSucc
119    if #gradInputSucc == 1 then
120       node.gradOutput = gradInputSucc[1]
121    elseif #gradInputSucc > 1 then
122       for k = 1, #gradInputSucc do
123          node.gradOutput = self:nestedAccTensor(node.gradOutput, gradInputSucc[k], k == 1)
124       end
125    end
126 end
127
128 ----------------------------------------------------------------------
129
130 -- Connect a sequence of modules
131 function DAG:connect(...)
132    self.sorted = nil
133    local prev
134    for _, nnm in pairs({...}) do
135       self:createNode(nnm)
136       if prev then
137          table.insert(self.node[nnm].pred, prev)
138          table.insert(self.node[prev].succ, nnm)
139       end
140       prev = nnm
141    end
142 end
143
144 function DAG:setInput(i)
145    self.sorted = nil
146    self.inputModules = i
147    self:nestedApply(
148       function(nnm)
149          assert(#self.node[nnm].succ > 0, 'Input modules must have outgoing edges.')
150          assert(#self.node[nnm].pred == 0, 'Input modules cannot have incoming edges.')
151       end,
152       self.inputModules
153    )
154 end
155
156 function DAG:setOutput(o)
157    self.sorted = nil
158    self.outputModules = o
159    self:nestedApply(
160       function(nnm)
161          assert(#self.node[nnm].pred > 0, 'Output module must have incoming edges.')
162          assert(#self.node[nnm].succ == 0, 'Output module cannot have outgoing edges.')
163       end,
164       self.outputModules
165    )
166 end
167
168 function DAG:print()
169    self:putInOrder()
170
171    for i, d in ipairs(self.sorted) do
172       print('#' .. i .. ' -> ' .. torch.type(d))
173    end
174 end
175
176 ----------------------------------------------------------------------
177
178 function DAG:saveDot(filename)
179    local file = (filename and io.open(filename, 'w')) or io.stdout
180
181    file:write('digraph {\n')
182
183    file:write('\n')
184
185    for nnmb, node in pairs(self.node) do
186       file:write(
187          '  '
188             .. node.index
189             .. ' [shape=box,label=\"' .. torch.type(nnmb) .. '\"]'
190             .. '\n'
191       )
192
193       for i, nnma in pairs(node.pred) do
194          local decoration = ''
195          if #node.pred > 1 then
196             -- decoration = ' [headlabel=\"' .. i .. '\"]'
197             decoration = ' [label=\"' .. i .. '\"]'
198          end
199          file:write(
200             '  '
201                .. self.node[nnma].index
202                .. ' -> '
203                .. self.node[nnmb].index
204                .. decoration
205                .. '\n'
206          )
207       end
208
209       file:write('\n')
210    end
211
212    file:write('}\n')
213
214 end
215
216 ----------------------------------------------------------------------
217
218 function DAG:updateOutput(input)
219    self:putInOrder()
220
221    self:nestedApply(
222       function(nnm, i)
223          local node = self.node[nnm]
224          node.input = i
225          self:rethrowErrors(nnm, node.index, 'updateOutput', i)
226       end,
227       self.inputModules,
228       input
229    )
230
231    for _, nnm in ipairs(self.sorted) do
232       local node = self.node[nnm]
233       if #node.pred > 0 then
234          local i
235          if #node.pred == 1 then
236             i = node.pred[1].output
237          elseif #node.pred > 1 then
238             i = {}
239             for k = 1, #node.pred do
240                i[k] = node.pred[k].output
241             end
242          end
243          node.input = i
244          self:rethrowErrors(nnm, node.index, 'updateOutput', i)
245       end
246    end
247
248    self.output = self:nestedApply(
249       function(m) return m.output end,
250       self.outputModules
251    )
252
253    return self.output
254 end
255
256 function DAG:updateGradInput(input, gradOutput)
257    assert(self.sorted, 'There has been a DAG structure change before a DAG:updateGradInput')
258
259    self:nestedApply(
260       function(nnm, go)
261          local node = self.node[nnm]
262          node.gradOutput = go
263          self:rethrowErrors(nnm, node.index, 'updateGradInput', node.input, go)
264       end,
265       self.outputModules, gradOutput
266    )
267
268    self:nestedApply(
269       function(nnm, i) self.node[nnm].input = i end,
270       self.inputModules, input
271    )
272
273    for _, node in pairs(self.node) do
274       node.gradInputSucc = {}
275    end
276
277    for k = #self.sorted, 1, -1 do
278       local nnm = self.sorted[k]
279       local node = self.node[nnm]
280       local pred = node.pred
281
282       if #node.gradInputSucc > 0 then
283          self:updateGradOutput(node)
284          self:rethrowErrors(nnm, node.index, 'updateGradInput', node.input, node.gradOutput)
285       end
286
287       -- We fill the gradInputSucc of our predecessors
288       if #pred == 1 then
289          table.insert(self.node[pred[1]].gradInputSucc, nnm.gradInput)
290       elseif #pred > 1 then
291          assert(torch.type(nnm.gradInput) == 'table',
292                 'Should have a table gradInput since it has multiple predecessors')
293          for n = 1, #pred do
294             table.insert(self.node[node.pred[n]].gradInputSucc, nnm.gradInput[n])
295          end
296       end
297    end
298
299    self.gradInput = self:nestedApply(function(m) return m.gradInput end, self.inputModules)
300
301    return self.gradInput
302 end
303
304 function DAG:accGradParameters(input, gradOutput, scale)
305    assert(self.sorted, 'There has been a DAG structure change before a DAG:accGradParameters')
306
307    self:nestedApply(
308       function(nnm, go) self.node[nnm].gradOutput = go end,
309       self.outputModules, gradOutput
310    )
311
312    self:nestedApply(
313       function(nnm, i) self.node[nnm].input = i end,
314       self.inputModules, input
315    )
316
317    for k = 1, #self.modules do
318       local nnm = self.modules[k]
319       local node = self.node[nnm]
320       self:rethrowErrors(nnm, k, 'accGradParameters', node.input, node.gradOutput, scale)
321    end
322 end
323
324 function DAG:clearState()
325    self.sorted = nil
326    for _, node in pairs(self.node) do
327       node.gradInputSucc = nil
328       node.input = nil
329       node.gradOutput = nil
330    end
331    return parent.clearState(self)
332 end