torch学习笔记3.2:实现自定义模块(cpu)

在使用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 文件说明

ReLU.lua

local ReLU = torch.class('nxn.ReLU')

function ReLU:__init(gpucompatible)

   self.gpucompatible=gpucompatible

   if self.gpucompatible then
      self.gradInput=torch.CudaTensor()
      self.output=torch.CudaTensor()
   else
      self.gradInput=torch.Tensor()
      self.output=torch.Tensor()
   end

   self.outputSave=self.output
   self.gradInputSave=self.gradInput

end

function ReLU:updateOutput(input)
   -- 调用cpp实现的ReLU函数
   return input.nxn.ReLU_updateOutput(self, input)
end

function ReLU:updateGradInput(input, gradOutput)
   -- 调用cpp实现的ReLU函数
   return input.nxn.ReLU_updateGradInput(self, input, gradOutput)
end

function ReLU:getDisposableTensors()
   return {self.output, self.gradInput, self.gradInputSave, self.outputSave}
end

ReLU.c

内容如下:

#ifndef TH_GENERIC_FILE
#define TH_GENERIC_FILE "generic/ReLU.c"
#else

static int nxn_(ReLU_updateOutput)(lua_State *L)
{
  printf("CPU version of ReLU updateOutput function\n");
  THTensor *input = luaT_checkudata(L, 2, torch_Tensor);
  THTensor *output = luaT_getfieldcheckudata(L, 1, "output", torch_Tensor);

  THTensor_(resizeAs)(output, input);

  TH_TENSOR_APPLY2(real, output, real, input,         \
                   *output_data = *input_data > 0 ? *input_data : 0;)

  return 1;
}

static int nxn_(ReLU_updateGradInput)(lua_State *L)
{
  printf("CPU version of ReLU updateGradInput function\n");
  THTensor *gradOutput = luaT_checkudata(L, 3, torch_Tensor);
  THTensor *output = luaT_getfieldcheckudata(L, 1, "output", torch_Tensor);
  THTensor *gradInput = luaT_getfieldcheckudata(L, 1, "gradInput", torch_Tensor);

  THTensor_(resizeAs)(gradInput, output);
  TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, output,     \
                   *gradInput_data = *gradOutput_data * (*output_data > 0 ? 1 : 0););
  return 1;
}

static const struct luaL_Reg nxn_(ReLU__) [] = {
  {"ReLU_updateOutput", nxn_(ReLU_updateOutput)},
  {"ReLU_updateGradInput", nxn_(ReLU_updateGradInput)},
  {NULL, NULL}
};

static void nxn_(ReLU_init)(lua_State *L)
{
  luaT_pushmetatable(L, torch_Tensor);
  luaT_registeratname(L, nxn_(ReLU__), "nxn");
  lua_pop(L,1);
}

#endif

init.c

在编译安装模块时CMakeLists.txt根据init.c找类文件:

#include "TH.h"
#include "luaT.h"

#define torch_(NAME) TH_CONCAT_3(torch_, Real, NAME)
#define torch_Tensor TH_CONCAT_STRING_3(torch.,Real,Tensor)
#define nxn_(NAME) TH_CONCAT_3(nxn_, Real, NAME)

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


LUA_EXTERNC DLL_EXPORT int luaopen_libnxn(lua_State *L);
// 把cpp实现编译到libnxn
int luaopen_libnxn(lua_State *L)
{
  lua_newtable(L);
  lua_pushvalue(L, -1);
  lua_setfield(L, LUA_GLOBALSINDEX, "nxn");

  nxn_FloatReLU_init(L);
  nxn_DoubleReLU_init(L);

  return 1;
}
  • 1
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
在 PyTorch 中,我们可以通过继承 `nn.Module` 类来实现自定义模块。下面是一个简单的例子,展示了如何实现一个简单的全连接层模块。 ```python import torch.nn as nn class MyLinear(nn.Module): def __init__(self, input_size, output_size): super(MyLinear, self).__init__() self.weight = nn.Parameter(torch.Tensor(output_size, input_size)) self.bias = nn.Parameter(torch.Tensor(output_size)) nn.init.xavier_uniform_(self.weight) nn.init.zeros_(self.bias) def forward(self, x): x = x.view(x.size(0), -1) # 将输入数据展平,以适应全连接层 out = torch.mm(x, self.weight.t()) + self.bias return out ``` 在上述代码中,我们定义了一个名为 `MyLinear` 的自定义模块,该模块继承自 `nn.Module` 类。我们在 `__init__` 方法中定义了模块的参数,包括权重和偏置,这些参数都是 `nn.Parameter` 类型,它们会自动被注册到模型参数列表中。我们还使用 `nn.init` 模块中的函数来初始化权重和偏置。在 `forward` 方法中,我们对输入数据进行展平操作,并使用 `torch.mm` 函数进行矩阵乘法运算,最后加上偏置即可。 使用自定义模块和内置模块的方式是一致的。例如,我们可以这样使用上述定义的自定义全连接层模块: ```python import torch # 创建一个 MyLinear 实例 linear = MyLinear(784, 10) # 随机生成一个大小为 (1, 784) 的张量 x = torch.randn(1, 784) # 使用自定义模块进行前向计算 out = linear(x) print(out.size()) # 输出:torch.Size([1, 10]) ``` 这里我们创建了一个输入大小为 784,输出大小为 10 的全连接层模块,并对一个大小为 (1, 784) 的随机输入数据进行前向计算,得到了大小为 (1, 10) 的输出。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值