torch学习笔记3.3：实现自定义模块(gpu)

在使用torch时，如果想自己实现一个层，则可以按照《torch学习笔记1：实现自定义层》 中的方法来实现。但是如果想要实现一个比较复杂的网络，往往需要自己实现多个层（或类），并且有时可能需要重写其他模块中已有的函数来达到自己的目的，如果还是在nn模块中添加，会比较混乱，并且不利于本地git仓库统一管理，这个时候，我们可以自己实现一个像nn一样的模块，在代码中使用时 require即可。

我们来实现一个名为nxn的自定义模块，以及它的cuda版本cunxn模块，其中包含一个自定义的Hello类（lua实现），ReLU类（分别用CPU和GPU实现）。

由于篇幅原因，这里把torch自定义模块的lua实现，cpu实现，gpu实现分别写一篇文章，本文介绍cpu实现的ReLU类。

1 总目录结构和 2 使用说明 在 《torch学习笔记3.1：实现自定义模块(lua)》

3 文件说明

这里介绍的都是cunxn文件夹里面的。

CMakeLists.txt

可以参考torch自带模块来写，主要是cuda文件的编译和链接，需要注意的部分内容如下：

......FIND_PACKAGE(CUDA 4.0 REQUIRED)SET(src-cuda init.cu)CUDA_ADD_LIBRARY(cunxn MODULE ${src-cuda})TARGET_LINK_LIBRARIES(cunxn luaT THC TH)IF(APPLE)  SET_TARGET_PROPERTIES(cunxn PROPERTIES    LINK_FLAGS "-undefined dynamic_lookup")ENDIF()### Torch packages supposes libraries prefix is "lib"SET_TARGET_PROPERTIES(cunxn PROPERTIES  PREFIX "lib"  IMPORT_PREFIX "lib")INSTALL(TARGETS cunxn  RUNTIME DESTINATION "${Torch_INSTALL_LUA_CPATH_SUBDIR}"  LIBRARY DESTINATION "${Torch_INSTALL_LUA_CPATH_SUBDIR}")SET(luasrc init.lua)INSTALL(  FILES  ${luasrc}  DESTINATION "${Torch_INSTALL_LUA_PATH_SUBDIR}/cunxn")ADD_TORCH_PACKAGE(cunxn "" "${luasrc}")

cunxn-scm-1.rockspec

其中的build部分和其他rockspec文件一样

package = "cunxn"version = "scm-1"source = {   url = "git://github.com/soumith/examplepackage.torch",   tag = "master"}dependencies = {   "torch >= 7.0",   "cunn",   "nn"}......

init.cu

同init.c的功能一样，编译时查找要编译的文件，以及生成libcunxn：

#include "luaT.h"#include "THC.h"#include "THLogAdd.h" /* DEBUG: WTF */#include <thrust/transform.h>#include <thrust/reduce.h>#include <thrust/transform_reduce.h>#include <thrust/functional.h>#include <thrust/device_ptr.h>#include "ReLU.cu"LUA_EXTERNC DLL_EXPORT int luaopen_libcunxn(lua_State *L);int luaopen_libcunxn(lua_State *L){  lua_newtable(L);  cunxn_ReLU_init(L);  return 1;}

init.lua

require "cutorch"require "nxn"require "libcunxn"

ReLU.cu

cuda实现的ReLU

struct reluupdateOutput_functor{  __host__ __device__ float operator()(const float& input) const  {    return input > 0 ? input : 0;  }};THCState* getCutorchState(lua_State* L){    lua_getglobal(L, "cutorch");    lua_getfield(L, -1, "getState");    lua_call(L, 0, 1);    THCState *state = (THCState*) lua_touserdata(L, -1);    lua_pop(L, 2);    return state;} static int cunxn_ReLU_updateOutput(lua_State *L){  printf("GPU version of ReLU updateOutput function\n");  THCState *state = getCutorchState(L);  THCudaTensor *input = (THCudaTensor*)luaT_checkudata(L, 2, "torch.CudaTensor");  THCudaTensor *output = (THCudaTensor*)luaT_getfieldcheckudata(L, 1, "output", "torch.CudaTensor");  long size = THCudaTensor_nElement(state, input);  input = THCudaTensor_newContiguous(state, input);  THCudaTensor_resizeAs(state, output, input);  thrust::device_ptr<float> output_data(THCudaTensor_data(state, output));  thrust::device_ptr<float> input_data(THCudaTensor_data(state, input));  thrust::transform(input_data, input_data+size, output_data, reluupdateOutput_functor());  THCudaTensor_free(state, input);  return 1;}struct reluupdateGradInput_functor{  __host__ __device__ float operator()(const float& output, const float& gradOutput) const  {    return gradOutput * (output > 0 ? 1 : 0);  }};static int cunxn_ReLU_updateGradInput(lua_State *L){  printf("GPU version of ReLU updateGradInput function\n");  THCState *state = getCutorchState(L);  THCudaTensor *output = (THCudaTensor*)luaT_getfieldcheckudata(L, 1, "output", "torch.CudaTensor");  THCudaTensor *gradOutput = (THCudaTensor*)luaT_checkudata(L, 3, "torch.CudaTensor");  THCudaTensor *gradInput = (THCudaTensor*)luaT_getfieldcheckudata(L, 1, "gradInput", "torch.CudaTensor");  long size = THCudaTensor_nElement(state, output);  gradOutput = THCudaTensor_newContiguous(state, gradOutput);  THCudaTensor_resizeAs(state, gradInput, output);  thrust::device_ptr<float> output_data(THCudaTensor_data(state, output));  thrust::device_ptr<float> gradOutput_data(THCudaTensor_data(state, gradOutput));  thrust::device_ptr<float> gradInput_data(THCudaTensor_data(state, gradInput));  thrust::transform(output_data, output_data+size, gradOutput_data, gradInput_data, reluupdateGradInput_functor());  THCudaTensor_free(state, gradOutput);  return 1;}static const struct luaL_Reg cunxn_ReLU__ [] = {  {"ReLU_updateOutput", cunxn_ReLU_updateOutput},  {"ReLU_updateGradInput", cunxn_ReLU_updateGradInput},  {NULL, NULL}};static void cunxn_ReLU_init(lua_State *L){  luaT_pushmetatable(L, "torch.CudaTensor");  luaT_registeratname(L, cunxn_ReLU__, "nxn");  lua_pop(L,1);}