当我们要实现自己的一些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'}"
即可测试。