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);
}