【pytorch源码赏析】nn.XX VS nn.functional.XX

来源:互联网 发布:大数据 图书 编辑:程序博客网 时间:2024/06/03 01:42

先来看代码

import torchimport torch.nn as nnimport torch.nn.functional as Fclass Net(nn.Module):    def __init__(self):        super(Net, self).__init__()        self.conv1_1 = nn.Conv2d(3, 64, 3, padding=1)        self.conv1_2 = nn.Conv2d(64, 3, 3, padding=1)        self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)    def forward(self, x):        x = F.relu(self.conv1_1(x))        x = self.pool1(F.relu(self.conv1_2(x)))        return x

定义网络的时候,有时候用nn.Conv2d,有时候用F.relu,两者有什么区别吗?

从上面代码可以看出,nn.Conv2d要先声明,像上面的8,9,10行,然后再调用,参与构建模型。而F.relu,直接使用,并没有声明的环节。

看源码。F.relu的源码,nn.Conv2d的源码

# F.reludef relu(input, inplace=False):    return _functions.thnn.Threshold.apply(input, 0, 0, inplace)

F.relu仅仅是一个函数,参数包括输入和计算所需参数,返回计算结果。它不能存储任何上下文信息。所有的function函数都从基类Function派生,实现forward和backward静态方法。而在forward和backward实现内部,调用了C的后端实现。比如下面这个MaxPool1d:

class MaxPool1d(Function):    @staticmethod    def forward(ctx, input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False):        ...        backend = type2backend[type(input)]        backend.SpatialDilatedMaxPooling_updateOutput(...)        ...        return output, indices    @staticmethod    def backward(ctx, grad_output, _indices_grad=None):        ...        grad_input = MaxPool1dBackward.apply(...)        return grad_input, None, None, None, None, None, None

nn.Conv2d是在F.conv2d外加的一层封装,从Module派生而来。Module是最接近用户的类,因为如果用户要实现自己的模型,也要从Module派生。Module提供了一种嵌套的模型定义方式,用户可以用这种简单的方式,定义很复杂的模型。关于Module最重要的一点是,它实现了__call__方法,在__call__方法里调用了forward()方法。所以,当执行self.Conv2d(x),它就在调用Conv2d里面的forward(),并返回计算结果。

nn.xx仅封装了一部分常用的计算模块,大部分还是在F.xx里面。如果需要自己用c实现operator(大部分情况不用),最后把接口提供给F.xx就可以。

另外,定义损失也是和nn.Conv2d一样的做法,先声明,再使用。