當我們要實現自己的一些idea時，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/NewClass.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'}"即可測試。