aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/torch/nn/Jacobian.lua
blob: 4f728b18c570fe19b999f294b10a09bea7909b8e (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
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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
nn.Jacobian = {}

function nn.Jacobian.backward(module, input, param, dparam)
   local doparam = 0
   if param then
      doparam = 1
   end
   param = param or input
   -- output deriv
   module:forward(input)
   local dout = module.output.new():resizeAs(module.output)
   -- 1D view
   local sdout = module.output.new(dout:storage(),1,dout:nElement())
   -- jacobian matrix to calculate
   local jacobian = torch.Tensor(param:nElement(),dout:nElement()):zero()

   for i=1,sdout:nElement() do
      dout:zero()
      sdout[i] = 1
      module:zeroGradParameters()
      local din = module:updateGradInput(input, dout)
      module:accGradParameters(input, dout)
      if doparam == 1 then
         jacobian:select(2,i):copy(dparam)
      else
         jacobian:select(2,i):copy(din)
      end
   end
   return jacobian
end

function nn.Jacobian.backwardUpdate(module, input, param)

   -- output deriv
   module:forward(input)
   local dout = module.output.new():resizeAs(module.output)
   -- 1D view
   local sdout = module.output.new(dout:storage(),1,dout:nElement())
   -- jacobian matrix to calculate
   local jacobian = torch.Tensor(param:nElement(),dout:nElement()):zero()

   -- original param
   local params = module:parameters()
   local origparams = {}
   for j=1,#params do
      table.insert(origparams, params[j]:clone())
   end

   for i=1,sdout:nElement() do
      for j=1,#params do
         params[j]:copy(origparams[j])
      end
      dout:zero()
      sdout[i] = 1
      module:updateGradInput(input, dout)
      module:accUpdateGradParameters(input, dout, 1)
      jacobian:select(2,i):copy(param)
   end

   for j=1,#params do
      params[j]:copy(origparams[j])
   end

   return jacobian
end

function nn.Jacobian.forward(module, input, param, perturbation)
   param = param or input
   -- perturbation amount
   perturbation = perturbation or 1e-6
   -- 1D view of input
   --local tst = param:storage()
   local sin = param.new(param):resize(param:nElement())--param.new(tst,1,tst:size())
   -- jacobian matrix to calculate
   local jacobian = torch.Tensor():resize(param:nElement(),module:forward(input):nElement())

   local outa = torch.Tensor(jacobian:size(2))
   local outb = torch.Tensor(jacobian:size(2))

   for i=1,sin:nElement() do
      local orig = sin[i]
      sin[i] = orig - perturbation
      outa:copy(module:forward(input))
      sin[i] = orig + perturbation
      outb:copy(module:forward(input))
      sin[i] = orig

      outb:add(-1,outa):div(2*perturbation)
      jacobian:select(1,i):copy(outb)
   end

   return jacobian
end

function nn.Jacobian.backwardDiagHessian(module, input, diagHessianParamName)
   -- Compute the second derivatives (diagonal Hessian elements)
   -- by backpropagation (using the code from hessian.lua).
   --
   -- This function computes the diagonal Hessian elements of the following function:
   --
   -- F(x_1, x_2, ..., x_n) = y_1^2/2 + y_2^2/2 + ... + y_m^2/2,
   --
   -- where
   -- x_1, ..., x_n are the input values and parameters of the given module,
   -- y_1, ..., y_m are the output values of the given module.
   --
   -- All x_i and y_i values are scalars here. In other words,
   -- x_1, ..., x_n denote the scalar elements of the module input tensor,
   --             the scalar elements of module.weight,
   --             and the scalar elements of module.bias;
   -- y_1, ..., y_m are the scalar elements of the module output tensor.
   --
   -- The diagonal Hessian elements of F are computed with respect to
   -- the module input values and parameters (x_1, .., x_n).
   --
   -- The function F is chosen for its convenient properties:
   --
   -- dF / dy_i = y_i,
   -- d^2F / dy_i^2 = 1.
   --
   -- In other words, the diagonal Hessian elements of F with respect
   -- to the module OUTPUT values (y_1, ... y_m) are equal to 1.
   --
   -- Because of that, computing the diagonal Hessian elements of F
   -- with respect to the module INPUT values and PARAMETERS (x_1, ..., x_n)
   -- can be done by calling updateDiagHessianInput() and accDiagHessianParameters()
   -- using a tensor of ones as diagHessianOutput.

   module:forward(input)
   local diagHessianOutput = module.output.new():resizeAs(module.output):fill(1)

   module.diagHessianWeight:zero()
   module.diagHessianBias:zero()
   module:updateDiagHessianInput(input, diagHessianOutput)
   module:accDiagHessianParameters(input, diagHessianOutput)

   return module[diagHessianParamName]
end

function nn.Jacobian.linearModuleDiagHessian(module, input, gradParamName)
   -- Compute the second derivatives (diagonal Hessian elements)
   -- from the first derivatives for the given module
   -- (without using the code from hessian.lua).
   --
   -- The given module is assumed to be linear with respect to its inputs and weights
   -- (like nn.Linear, nn.SpatialConvolution, etc.)
   --
   -- This function computes the diagonal Hessian elements of the following function:
   --
   -- F(x_1, x_2, ..., x_n) = y_1^2/2 + y_2^2/2 + ... + y_m^2/2.
   --
   -- (See the the comment for nn.Jacobian.backwardDiagHessian() for explanation.)
   --
   -- The first derivatives of F with respect to
   -- the module inputs and parameters (x_1, ..., x_n) are:
   --
   -- dF / dx_i = \sum_k (dF / dy_k) (dy_k / dx_i).
   --
   -- The second derivatives are:
   --
   -- d^2F / dx_i = \sum_k [(d^2F / dy_k^2) (dy_k / dx_i)^2 + (dF / dy_k) (d^2y_k / dx_i^2)].
   --
   -- The second derivatives of F with respect to the module outputs (y_1, ..., y_m)
   -- are equal to 1, so:
   --
   -- d^2F / dx_i = \sum_k [(dy_k / dx_i)^2 + (dF / dy_k) (d^2y_k / dx_i^2)].
   --
   -- Assuming the linearity of module outputs (y_1, ..., y_m)
   -- with respect to module inputs and parameters (x_1, ..., x_n),
   -- we have (d^2y_k / dx_i^2) = 0,
   -- and the expression finally becomes:
   --
   -- d^2F / dx_i = \sum_k (dy_k / dx_i)^2.
   --
   -- The first derivatives (dy_k / dx_i) are computed by normal backpropagation,
   -- using updateGradInput() and accGradParameters().

   local gradParam = module[gradParamName]

   local diagHessian = gradParam.new():resize(gradParam:nElement()):zero()

   module:forward(input)
   local gradOutput = module.output.new():resizeAs(module.output)
   local gradOutput1D = gradOutput:view(gradOutput:nElement())

   for i=1,gradOutput:nElement() do
      gradOutput1D:zero()
      gradOutput1D[i] = 1
      module.gradWeight:zero()
      if module.bias then
         module.gradBias:zero()
      end
      module:updateGradInput(input, gradOutput)
      module:accGradParameters(input, gradOutput)
      diagHessian:addcmul(gradParam, gradParam)
   end

   return diagHessian
end

function nn.Jacobian.forwardUpdate(module, input, param, perturbation)
   -- perturbation amount
   perturbation = perturbation or 1e-6
   -- 1D view of input
   --local tst = param:storage()
   local sin =  param.new(param):resize(param:nElement())--param.new(tst,1,tst:size())
   -- jacobian matrix to calculate
   local jacobian = torch.Tensor():resize(param:nElement(),module:forward(input):nElement())

   local outa = torch.Tensor(jacobian:size(2))
   local outb = torch.Tensor(jacobian:size(2))

   for i=1,sin:nElement() do
      local orig = sin[i]
      sin[i] = orig - perturbation
      outa:copy(module:forward(input))
      sin[i] = orig + perturbation
      outb:copy(module:forward(input))
      sin[i] = orig

      outb:add(-1,outa):div(2*perturbation)
      jacobian:select(1,i):copy(outb)
      jacobian:select(1,i):mul(-1)
      jacobian:select(1,i):add(sin[i])
   end
   return jacobian
end

function nn.Jacobian.testJacobian(module, input, minval, maxval, perturbation)
   minval = minval or -2
   maxval = maxval or 2
   local inrange = maxval - minval
   input:copy(torch.rand(input:nElement()):mul(inrange):add(minval))
   local jac_fprop = nn.Jacobian.forward(module, input, input, perturbation)
   local jac_bprop = nn.Jacobian.backward(module, input)
   local error = jac_fprop-jac_bprop
   return error:abs():max()
end

function nn.Jacobian.testJacobianParameters(module, input, param, dparam, minval, maxval, perturbation)
   minval = minval or -2
   maxval = maxval or 2
   local inrange = maxval - minval
   input:copy(torch.rand(input:nElement()):mul(inrange):add(minval))
   param:copy(torch.rand(param:nElement()):mul(inrange):add(minval))
   local jac_bprop = nn.Jacobian.backward(module, input, param, dparam)
   local jac_fprop = nn.Jacobian.forward(module, input, param, perturbation)
   local error = jac_fprop - jac_bprop
   return error:abs():max()
end

function nn.Jacobian.testJacobianUpdateParameters(module, input, param, minval, maxval, perturbation)
   minval = minval or -2
   maxval = maxval or 2
   local inrange = maxval - minval
   input:copy(torch.rand(input:nElement()):mul(inrange):add(minval))
   param:copy(torch.rand(param:nElement()):mul(inrange):add(minval))
   local params_bprop = nn.Jacobian.backwardUpdate(module, input, param)
   local params_fprop = nn.Jacobian.forwardUpdate(module, input, param, perturbation)

   local error = params_fprop - params_bprop
   return error:abs():max()
end

function nn.Jacobian.testDiagHessian(module, input, gradParamName, diagHessianParamName, minval, maxval)
   -- Compute the diagonal Hessian elements for the same function in two different ways,
   -- then compare the results and return the difference.

   minval = minval or -2
   maxval = maxval or 2
   local inrange = maxval - minval
   input:copy(torch.rand(input:nElement()):mul(inrange):add(minval))
   module:initDiagHessianParameters()
   local h_bprop = nn.Jacobian.backwardDiagHessian(module, input, diagHessianParamName)
   local h_linearmodule = nn.Jacobian.linearModuleDiagHessian(module, input, gradParamName)
   local error = h_bprop - h_linearmodule
   return error:abs():max()
end

function nn.Jacobian.testDiagHessianInput(module, input, minval, maxval)
   return nn.Jacobian.testDiagHessian(module, input, 'gradInput', 'diagHessianInput', minval, maxval)
end

function nn.Jacobian.testDiagHessianWeight(module, input, minval, maxval)
   return nn.Jacobian.testDiagHessian(module, input, 'gradWeight', 'diagHessianWeight', minval, maxval)
end

function nn.Jacobian.testDiagHessianBias(module, input, minval, maxval)
   return nn.Jacobian.testDiagHessian(module, input, 'gradBias', 'diagHessianBias', minval, maxval)
end

function nn.Jacobian.testIO(module,input, minval, maxval)
   minval = minval or -2
   maxval = maxval or 2
   local inrange = maxval - minval
   local inputclone = input:clone()

   -- run module
   module:forward(input)
   local go = module.output:clone():copy(torch.rand(module.output:nElement()):mul(inrange):add(minval))
   local goclone = go:clone()
   module:zeroGradParameters()
   module:updateGradInput(input,go)
   module:accGradParameters(input,go)

   local fo = module.output:clone()
   local bo = module.gradInput:clone()

   -- write module
   local filename = os.tmpname()
   local f = torch.DiskFile(filename, 'w'):binary()
   -- call clearState and check that it returns itself
   assert(module == module:clearState(),'clearState did not return self')
   f:writeObject(module)
   f:close()
   -- read module
   local m = torch.DiskFile(filename):binary():readObject()
   m:forward(inputclone)
   m:zeroGradParameters()
   m:updateGradInput(inputclone,goclone)
   m:accGradParameters(inputclone,goclone)
   -- cleanup
   os.remove(filename)

   local fo2 = m.output:clone()
   local bo2 = m.gradInput:clone()

   local errf = fo - fo2
   local errb = bo - bo2
   return errf:abs():max(), errb:numel() == 0 and 0 or errb:abs():max()
end

function nn.Jacobian.testAllUpdate(module, input, weight, gradWeight)
   local gradOutput
   local lr = torch.uniform(0.1, 1)
   local errors = {}

   -- accGradParameters
   local maccgp = module:clone()
   local weightc = maccgp[weight]:clone()
   maccgp:forward(input)
   gradOutput = torch.rand(maccgp.output:size())
   maccgp:zeroGradParameters()
   maccgp:updateGradInput(input, gradOutput)
   maccgp:accGradParameters(input, gradOutput)
   maccgp:updateParameters(lr)
   errors["accGradParameters"] = (weightc-maccgp[gradWeight]*lr-maccgp[weight]):norm()

   -- accUpdateGradParameters
   local maccugp = module:clone()
   maccugp:forward(input)
   maccugp:updateGradInput(input, gradOutput)
   maccugp:accUpdateGradParameters(input, gradOutput, lr)
   errors["accUpdateGradParameters"] = (maccugp[weight]-maccgp[weight]):norm()

   -- shared, accGradParameters
   local macsh1 = module:clone()
   local macsh2 = module:clone()
   macsh2:share(macsh1, weight)
   macsh1:forward(input)
   macsh2:forward(input)
   macsh1:zeroGradParameters()
   macsh2:zeroGradParameters()
   macsh1:updateGradInput(input, gradOutput)
   macsh2:updateGradInput(input, gradOutput)
   macsh1:accGradParameters(input, gradOutput)
   macsh2:accGradParameters(input, gradOutput)
   macsh1:updateParameters(lr)
   macsh2:updateParameters(lr)
   local err = (weightc-maccgp[gradWeight]*(lr*2)-macsh1[weight]):norm()
   err = err + (weightc-maccgp[gradWeight]*(lr*2)-macsh2[weight]):norm()
   errors["accGradParameters [shared]"] = err

   -- shared, accUpdateGradParameters
   local macshu1 = module:clone()
   local macshu2 = module:clone()
   macshu2:share(macshu1, weight)
   macshu1:forward(input)
   macshu2:forward(input)
   macshu1:updateGradInput(input, gradOutput)
   macshu2:updateGradInput(input, gradOutput)
   macshu1:accUpdateGradParameters(input, gradOutput, lr)
   macshu2:accUpdateGradParameters(input, gradOutput, lr)
   err = (weightc-maccgp[gradWeight]*(lr*2)-macshu1[weight]):norm()
   err = err + (weightc-maccgp[gradWeight]*(lr*2)-macshu2[weight]):norm()
   errors["accUpdateGradParameters [shared]"] = err

   return errors
end