扩展 torch.nn

nn 包含两种接口 - modules和他们的functional版本。通过这两个接口,你都可以扩展nn。但是我们建议,在扩展layer的时候,使用modules, 因为modules保存着参数和buffer。如果不需要参数的话,那么建议使用functional(激活函数,pooling,这些都不需要参数)。

增加一个operation的 functional版本已经在上面一节介绍完毕。

增加一个模块(module)。 由于nn重度使用autograd。所以,添加一个新module需要实现一个 用来执行 计算 和 计算梯度 的Function。从现在开始,假定我们想要实现一个Linear module,记得之前我们已经实现了一个Linear Funciton。 只需要很少的代码就可以完成这个工作。 现在,我们需要实现两个方法:

  • init (optional) - 输入参数,例如kernel sizes, numbers of features, 等等。同时初始化 parameters和buffers。
  • forward() - 实例化一个执行operation的Function,使用它执行operation。和functional wrapper(上面实现的那个简单的wrapper)十分类似。

Linear module实现代码:

class Linear(nn.Module): def init(self, inputfeatures, output_features, bias=True): self.input_features = input_features self.output_features = output_features # nn.Parameter is a special kind of Variable, that will get # automatically registered as Module's parameter once it's assigned # as an attribute. Parameters and buffers need to be registered, or # they won't appear in .parameters() (doesn't apply to buffers), and # won't be converted when e.g. .cuda() is called. You can use # .register_buffer() to register buffers. # nn.Parameters can never be volatile and, different than Variables, # they require gradients by default. self.weight = nn.Parameter(torch.Tensor(input_features, output_features)) if bias: self.bias = nn.Parameter(torch.Tensor(output_features)) else: # You should always register all possible parameters, but the # optional ones can be None if you want. self.register_parameter('bias', None) # Not a very smart way to initialize weights self.weight.data.uniform(-0.1, 0.1) if bias is not None: self.bias.data.uniform_(-0.1, 0.1) def forward(self, input): # See the autograd section for explanation of what happens here. return Linear()(input, self.weight, self.bias) #注意这个Linear是之前实现过的Linear

results matching ""

    No results matching ""