Torch 中添加自己的 nn Modules：以添加 Dropout、 Triplet Loss 为例

Preface

因为要复现前面阅读的一篇论文：《论文笔记：Deep Relative Distance Learning: Tell the Difference Between Similar Vehicles》 中提到的用来区分相似图像的两个损失函数：Triplet Loss、Coupled Cluster Loss 。上面的那篇论文没有提供源代码，因此得自己去写这两个损失函数模块，然后 require '***' 到 Torch 中。

这里学习一下该怎么样在 Torch 中加自己的模块层。Torch 有了官方教程：http://torch.ch/docs/developer-docs.html，中科院自动化所的博士@beanfrog 也写过一篇专栏：《如何在torch中添加新的层？》，下面是阅读及总结。

另，我还没到自己写 C/CUDA 代码的程度上，仅仅在用 lua、Tensor 实现的很浅的程度上，理解的不对请多多批评。

Torch 中如何实现自己的 nn 模块

首先，Torch 中实现自己的 nn modules 有两种层次：
1. 如果要实现的目标比较简单，对性能要求不是那么高，那么几行简单的 Lua 代码就可以实现；
2. 如果对计算性能要求较高，或者你想在用 CPU、GPU 计算时做特殊的处理、优化，就需要写 C/CUDA 层次的代码。

Modules 包含两种状态变量（state variables）：output 以及 gradInput。先回顾一下，一个 Module 需要去实现一些基础的函数。
第一个是如下的函数：

[output] forward(input)

这个函数是输入一个数据，计算模型相应的输出。通常来说，这里的 input 以及 output 都是 Tensors 。但是，一些特殊的 sub-classes，如 table layers 输入输出可能并不是 Tensors，而是其他什么东西。这个具体的需要去查阅每个 Module 的具体说明。
在 forward() 之后，变量 output 的值就需要更新到一个新的值。
但是，并不建议这么去 override 这个 forward() 函数。相反的，我们需要去实现 updateOutput(input) 这个函数。抽象父类 Module 中的 forward(input) 函数，会去调用 updateOutput(input)。


第二个需要去实现的函数是：

[gradInput] backward(input, gradOutput)

这个函数根据输入的 input，实现相应的 back-propagation 过程。一般来说，当调用这个函数时，应该是已经执行了 forward(input) 这个函数。

If you do not respect this rule, backend() will compute incorrect gradients.

gradOutput、gradInput都是 Tensors，同上，有些并不是，如 table 类，所以需要查每一个 Module 具体的说明。

一个 back-propagation 过程在给定的 gradOutput（Module 关于输出 output 的求导） 之上，包括了两个梯度计算。

这个函数调用了两个函数：

• 调用了 updateGradInput(input, gradOutput)
• 调用了 accGradParameters(input, gradOutput)

同样也不建议直接去重写 backward(input, gradOutput) 函数，更好的方式是去重写 updateGradInput(input, gradOutput) 、 accGradParameters(input, gradOutput)这两个函数。


好了，总结一下。
当定义一个新的 module 的时候，需要我们必须实现两种操作：forward、backward 。但并不建议去直接 override 这两个函数，而是去重写下面两个函数。第一个就是：

[output] updateOutput(input)

当定义一个 Module 的时候，上面的这个函数需要被重写。用当前的模型（模型里的参数不变），输入一个 input 值，得到输出。这个函数的返回值存储在 output 变量中。

第二个就是：

[gradInput] updateGradInput(input, gradOutput)

这个函数在新定义 Module 的时候也需要被重写，计算 Module 关于输入 input 的导数，函数的返回值为 gradInput，同时，变量 gradInput 的值需要进行更新。


接着要注意了，当定义一个新的 Module 时，如果这个 Module 有需要去训练的参数时（如激活函数层，就没有需要去训练的参数），下面的这个函数需要被 override 。

accGradParameters(input, gradOutput)

这是计算一个 Module 相对于权重的导数，许多 Module 不需要执行这一步，因为 没有权重参数，如激活函数层。


将这些 累积（accumulation） 归零（Zeroing）可以通过函数：zeroGradParameters() 来实现，根据 累积 更新这些参数可以通过函数：updateParameters() 来实现。


还可以定义函数：reset() ，这个函数定义怎么样去重置训练参数，如在训练前初始化参数。


如果你想使用 optim package ，modules 还提供了其他你可以去自定义的函数。


下面是一个定义新类的模板，每当我们要去定义一个新的 nn Module 的时候，我们只需要去填补下面函数体。

local NewClass, Parent = torch.class('nn.NewClass', 'nn.Module')

function NewClass:__init()
   Parent.__init(self)
   end

function NewClass:updateOutput(input)
end

function NewClass:updateGradInput(input, gradOutput)
end

function NewClass:accGradParameters(input, gradOutput)
end

function NewClass:reset()
end

要注意的是，当定义 __init() 这个 “构造函数（constructor）” 后，我们都首先要去调用这个 “构造函数（constructor）”，进行初始化操作。如我们常用的 nn.Linear(3, 4)，其表示 3 个输入，4 个输出，其初始化函数形式为 function Linear:__init(inputSize, outputSize, bias)， 其中 bias 表示偏置项，默认没有。

updateOutput(input)，前向传播时调用该函数，计算 $Y = F(X)$，$X$ 为输入，$Y$ 为输出。

updateGradInput(input, gradOutput)，反向传播是调用该函数，计算损失 $loss$ 相对于输入 $X$ 的偏导：$\frac{\partial E}{\partial X} = \frac{\partial E}{\partial Y} \times \frac{\partial Y}{\partial X}$，其中，$\frac{\partial E}{\partial Y}$ 为 $loss$ 相对于输出 $Y$ 的偏导。

accGradParameters(input, gradOutput)，当 模块（层） 中没有需要学习的 权重（参数） 时， 不需要写这个函数（如 nn.Dropout、nn.ReLU 等）。该函数计算输出 $loss$ 对 $weight$、$bias$ （如果有的话）的偏导：

$$\frac{\partial E}{\partial W} = \frac{\partial E}{\partial Y} \times \frac{\partial Y}{\partial W} \\ \frac{\partial E}{\partial B} = \frac{\partial E}{\partial Y} \times \frac{\partial Y}{\partial B}$$

以上函数的具体实现，若能通过 Tensor 自带的计算完成，则直接一个 Lua 文件即可，因为 Tensor 的运算自动支持 CPU 和 GPU，故该层也能直接支持 CPU 和 GPU了。

实现自己的 Dropout 层

Dropout 模块代码

所谓的 Dropout ，是指随机的将网络中的神经元归零，具体的可以去看 Hinton 的 JMLR 的这篇论文：《Dropout - A Simple Way to Prevent Neural Networks from Overfitting》，已经实验证实，这种策略可以有效的防止过拟合。

那么 Dropout 该如何实现呢？

local Dropout, Parent = torch.class('nn.Dropout', 'nn.Module')

function Dropout:__init(p)
   Parent.__init(self)
   self.p = p or 0.5
   if self.p >= 1 or self.p < 0 then
      error('<Dropout> illegal percentage, must be 0 <= p < 1')
   end
   self.noise = torch.Tensor()
end

function Dropout:updateOutput(input)
   self.output:resizeAs(input):copy(input)
   self.noise:resizeAs(input)
   self.noise:bernoulli(1-self.p)
   self.output:cmul(self.noise)
   return self.output
end

function Dropout:updateGradInput(input, gradOutput)
   self.gradInput:resizeAs(gradOutput):copy(gradOutput)
   self.gradInput:cmul(self.noise) -- simply mask the gradients with the noise vector
   return self.gradInput
end

这里还有一篇添加自己的 ReLU 层的例子，可以参考：http://blog.youkuaiyun.com/lanran2/article/details/50494570

Jacobian 矩阵检验

Jacobian 矩阵 是一阶偏导数以一定方式排列成的矩阵，类似于多元函数的导数。
假设 $F: R_n \rightarrow R_m$ 是一个从欧式 $n$ 维空间转换到欧式 $m$ 维空间的函数。这个函数由 $m$ 个实函数组成：

$$y_1(x_1, ..., x_n) \\ ... \\ y_m(x_1, ..., x_n)$$
这些函数的偏导数（如果存在）可以组成一个 $m$ 行 $n$ 列的矩阵，这就是所谓的 Jacobian 矩阵：

$$\begin{bmatrix} \frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_1}{\partial x_n} \\ \vdots & \ddots & \vdots \\ \frac{\partial y_m}{\partial x_1} & \cdots & \frac{\partial y_m}{\partial x_n} \end{bmatrix}$$
此矩阵表示为： ${J_F}({x_1}, \ldots ,{x_n})$ 或者 $\frac{{\partial ({y_1}, \ldots ,{y_m})}}{{\partial ({x_1}, \ldots ,{x_n})}}$

当实现完自己的 module 后，最好还需要测试验证。这可以用 nn 中提供的 Jacobian 类来检验：

-- parameters
local precision = 1e-5
local jac = nn.Jacobian

-- define inputs and module
local ini = math.random(10, 20)
local inj = math.random(10, 20)
local ink = math.random(10, 20)
local percentage = 0.5
local input = torch.Tensor(ini, inj, ink):zero()
local module = nn.Dropout(percentage)

-- test backprop, with  Jacobian
local err = jac.testJacobian(module, input)
print('==> error: ' .. err)
if err < precision then
    print('==> module OK')
else
    print('==> err too large, incorrect implementation')
end

上部分检验代码中，用了函数：nn.Jacobian.testJacobian()，这个函数的代码在这里：https://github.com/torch/nn/blob/master/Jacobian.lua#L240，翻进去看代码：

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

实现 Triplet Loss

Triplet Loss 示意图及其 Loss Function

Triplet Loss 的输入是“三元”的：$\{ <x^a,x^p,x^n> \}$，其中，$x^a$ 与 $x^p$ 属于正样本，$x^n$ 属于负样本。通过训练，使得 $x^a$ 与 $x^p$ 之间“拉的更近”，$x^a$ 与 $x^n$ 之间“推的更远”。示意图如下：

其损失函数为：

$$L = \sum^N{max\{ \left\| f(x^a) - f(x^p) \right\|^2_2 + \alpha - \left\| f(x^a) - f(x^n) \right\|^2_2,\ 0 \}}$$

Triplet Loss 模块 Torch 实现

Github 上这里有人实现了 Triplet Loss：https://github.com/Atcold/torch-TripletEmbedding，具体的实现很简单，就写了一个 TripletEmbedding.lua 文件：

local TripletEmbeddingCriterion, parent = torch.class('nn.TripletEmbeddingCriterion', 'nn.Criterion')

这是 Triplet Loss 类的实现声明，一般自定义 Loss Function 继承自 nn.Criterion.

function TripletEmbeddingCriterion:__init(alpha)
   parent.__init(self)
   self.alpha = alpha or 0.2
   self.Li = torch.Tensor()
   self.gradInput = {}
end

这是 Triplet Loss 的初始化函数，当使用 Triplet Loss 的时候，这个函数首先被调用，进行初始化。

• self.alpha 是 Triplet Loss 公式中的 $\alpha$，默认为 $0.2$
• self.Li 表示每一个计算得到的 $Loss$，先声明为 torch.Tensor()，大小维数、类型都待定
• self.gradInput 表示输入，输入为 table: {}，其实是为 {aImgs, pImgs, nImgs}
  
  

function TripletEmbeddingCriterion:updateOutput(input)
   local a = input[1] -- ancor
   local p = input[2] -- positive
   local n = input[3] -- negative
   local N = a:size(1) 
   self.Li = torch.max(torch.cat(torch.Tensor(N):zero():type(torch.type(a)) , (a - p):norm(2,2):pow(2) -  (a - n):norm(2,2):pow(2) + self.alpha, 2), 2)
   self.output = self.Li:sum() / N
   return self.output
end

这是需要我们自己去实现的 forward 过程，但是因为不建议去直接 override forward 函数，而是去实现 updateOutput() 函数。

输入的 input 实际上为 table，{aImgs, pImgs, nImgs}。所以创建 3 个 local 变量：local a、local p、local n；

local N 为 input 中传递过来的三元组的个数；

torch.max(0, x) 函数很好理解；

torch.Tensor(N):zero():type(torch.type(a))，创建一列 N 个 0 元素 Tensor，类型与 a 的一致；

(a - p):norm(2,2):pow(2)，因为若 a、p 的 Size 是 $N \times 1024$ ，$1024$ 是最后输出的特征维数（我选择的是 $1024$ 维），所以这句话(a - p):norm(2,2)求的是在第二个维度（$1024$）上的 2 范数，并 :pow(2)一下，即平方一下；

self.output = self.Li:sum() / N，求一个均值；

return self.output，返回结果。

function TripletEmbeddingCriterion:updateGradInput(input)
   local a = input[1] -- ancor
   local p = input[2] -- positive
   local n = input[3] -- negative
   local N = a:size(1)

   self.gradInput[1] = (n - p):cmul(self.Li:gt(0):repeatTensor(a:size(2),1):t():type(a:type()) * 2/N)
   self.gradInput[2] = (p - a):cmul(self.Li:gt(0):repeatTensor(a:size(2),1):t():type(a:type()) * 2/N)
   self.gradInput[3] = (a - n):cmul(self.Li:gt(0):repeatTensor(a:size(2),1):t():type(a:type()) * 2/N)

   return self.gradInput
end

这是反向传播过程需要实现了 updateGradInput函数，因为 Triplet Loss 的损失函数，对三个输入的求导分别为：

这是截取的 2015 年 CVPR 的一篇论文：《Simultaneous feature learning and hash coding with deep neural networks》上的公式。上面函数体里面就实现的是上图中 公式（3）的内容。

Reference

最后，在贴出几个有价值的链接，可供参考：

1. Google Groups 中的讨论：Custom softmax loss function
2. Oxford Machine Learning 教程：Implementing your own layer
3. http://code.madbits.com/wiki/doku.php?id=tutorial_morestuff
4. Stackoverflow 上的讨论：Add my custom loss function to torch
5. Github 上 Torch 的 Modules 文档，实现一个自己的 Modules，文档得看看