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

先来看代码

import torch
import torch.nn as nn
import torch.nn.functional as F

class 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.relu
def 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一样的做法，先声明，再使用。