torch中实现自定义层

当我们要实现自己的一些idea时，torch自带的模块和函数已经不能满足，我们需要自己实现层（或者类），一般的做法是把自定义层加入到已有的torch模块中。

torch中实现自定义层

lua实现

如果自定义层的功能可以通过调用torch中已有的函数实现，那就只需要用lua实现，torch的文档中也提供了简单的说明。
现在我们来实现一个NewClass：

• 在torch目录下（torch/extra/nn/）创建文件NewClass.lua

• 参考nn中其他lua文件的结构写好模板，在对应的函数中实现想要的功能

--创建新类，从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

• 在nn的init.lua中末尾添加一句

require('nn.NewClass')

• 重新安装nn模块

cd torch/extra/nn/
luarocks make rocks/nn-scm-1.rockspec

• 安装成功后，在自己的代码中使用自定义的类了

require 'nn'
...
nn.NewClass()
...

CPU实现

如果通过torch的函数不能实现出需要的功能，那么需要自己写C程序实现核心功能，然后在NewClass.lua中调用。

• 在torch/extra/nn/lib/THNN/generic/目录下创建文件NewClass.c

• 参考nn中已有的实现，在函数中实现需要的功能

...
void THNN_(NewClass_updateOutput)(
          THNNState *state,
          THTensor *input,
          THTensor *output)
{

}

void THNN_(NewClass_updateGradInput)(
          THNNState *state,
          THTensor *input,
          THTensor *gradOutput,
          THTensor *gradInput)
{

}
...

• 声明已实现的函数，在torch/extra/nn/lib/THNN/generic/THNN.h中添加

...
TH_API void THNN_(NewClass_updateOutput)(
          THNNState *state,            
          THTensor *input,             
          THTensor *output);           
TH_API void THNN_(NewClass_updateGradInput)(
          THNNState *state,            
          THTensor *input,            
          THTensor *gradOutput,      
          THTensor *gradInput);    
...    

• 添加include，在torch/extra/nn/lib/THNN/init.c中，添加

#include "generic/SpatialConvolution.c"
#include "THGenerateFloatTypes.h"

• 在NewClass.lua中调用CPU版本的函数

...
function NewClass:updateOutput(input)
   input.THNN.NewClass_updateOutput(
      input:cdata(),
      self.output:cdata()
   )
   return self.output
end

function NewClass:updateGradInput(input, gradOutput)

   if self.gradInput then
      input.THNN.NewClass_updateGradInput(
         input:cdata(),
         self.gradInput:cdata(),
         gradOutput:cdata()
      )
      return self.gradInput
   end
end
...

• 重新编译安装nn

cd torch/extra/nn/
luarocks make rocks/nn-scm-1.rockspec

• 安装成功后，在自己的代码中使用自定义的类了

require 'nn'
...
nn.NewClass()
...

Cuda实现

如果想要进一步提升运算效率，需要自己写一个Cuda版本的程序。

• 在torch/extra/cunn/lib/THCUNN/目录下创建文件NewClass.cu

• 参考cunn中已有的函数，实现函数功能

...
void THNN_CudaNewClass_updateOutput(THCState *state, THCudaTensor *input, THCudaTensor *output)
{

}

void THNN_CudaNewClass_updateGradInput(THCState *state, THCudaTensor *input, THCudaTensor *gradOutput, THCudaTensor *gradInput)
{

}
...

• 声明函数，在torch/extra/cunn/lib/THCUNN/THCUNN.h中添加：

TH_API void THNN_CudaNewClass_updateOutput(
          THCState *state,
          THCudaTensor *input,
          THCudaTensor *output);
TH_API void THNN_CudaNewClass_updateGradInput(
          THCState *state,
          THCudaTensor *input,
          THCudaTensor *gradOutput,
          THCudaTensor *gradInput);

• 在NewClass.lua中调用GPU版本的函数，和CPU版本一样，都通过THNN调用

• 重新编译安装cunn

cd torch/extra/cunn/
luarocks make rocks/cunn-scm-1.rockspec

• 安装成功后，在自己的代码中使用自定义的类了

require 'cunn'
...
nn.NewClass()
...

测试

在torch/extra/nn/test.lua和torch/extra/cunn/test.lua中添加测试代码，可以用来测试NewClass的输出是否正确，具体可参考已有的测试代码。
添加好后，运行th -lnn -e "nn.test{'NewClass'}"即可测试。