PyTorch 中torch.nn. 与torch.nn.function.的区别

今天我在定义一个卷积网络的过程中，发现nn.Conv2d与nn.functional.conv2d看似都是相同的。

1
2
3
4
5
class DogCat_Net(nn.Module):
    def __init__(self):
        super(Dog_Cat,self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size = 7,dilation = 2)
        self.conv1 = nn.functional.conv2d(3, 16, kernel_size = 7,dilation = 2)

类似的还有很多，在卷积，激活，池化等操作中，nn.Conv2d与nn.functional.conv2d均很相似。

通过查看PyTorch的文档发现，例如torch.nn.Conv2d是一个类

torch.nn.functional.conv2d则更像一个函数：

简而言之：

  • nn.Conv2d可以理解为一个Class类，需要继承nn.Module类
  • nn.functional.conv2d可以理解为一个纯函数，由def function(input)定义

对于nn中定义的类，可以提取变化的学习参数。而nn.functional中的是函数，是一个固定的运算公式。

由于在深度学习中会有很多权重是不断变化的，所以需要采用类的方式，以确保参数改变后，仍能正确进行运算。

所以一般在建立网络模型时，在__init__函数中，通常使用nn中的类；而在forward函数中，通常使用nn.functional中的函数。

例如以下网络模型：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Dog_Cat(nn.Module):
    def __init__(self):
        super(Dog_Cat,self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size = 7,dilation = 2)
        self.conv2 = nn.Conv2d(16, 32 ,kernel_size = 5)
        self.dropout = nn.Dropout2d()
        self.fc1 = nn.Linear(55696,1000)
        self.fc2 = nn.Linear(1000,50)
        self.fc3 = nn.Linear(50,2)
    def forward(self,x):
        x = F.relu(F.max_pool2d(self.conv1(x),2))
        x = F.relu(F.max_pool2d(self.dropout(self.conv2(x)),2))
        x = x.view(-1,55696)
        x = F.relu(self.fc1(x))
        x = F.dropout(x)
        x = F.relu(self.fc2(x))
        x = F.dropout(x)
        x = F.log_softmax(self.fc3(x),dim=1)
        return x

当网络模型使用了Dropout层时，其只需要在模型训练的时候进行使用，在模型测试时并不会使用它。

当Dropout层是使用nn.Dropout定义时，在模型训练时，可以通过model.train()开启它，在模型测试时，可以通过model.eval()关闭它。

但当Dropout层是用nn.functional.dropout定义的，那么在使用model.eval()后，也没办法关闭Dropout层。

以上提到了model.train()与model.eval()，那么它们的具体作用是什么呢？为什么要在模型训练与测试的时候加上这两段代码呢？

model.train()

如果模型中有BN层(Batch Normalization）和Dropout，需要在训练时添加model.train()。model.train()是保证BN层能够用到每一批数据的均值和方差。对于Dropout，model.train()是随机取一部分网络连接来训练更新参数。

model.eval()

如果模型中有BN层(Batch Normalization）和Dropout，在测试时添加model.eval()。model.eval()是保证BN层能够用全部训练数据的均值和方差，即测试过程中要保证BN层的均值和方差不变。对于Dropout，model.eval()是利用到了所有网络连接，即不进行随机舍弃神经元。

训练完train样本后，生成的模型model要用来测试样本。在model(test)之前，需要加上model.eval()，框架会自动把BN和Dropout固定住，不会取平均，而是用训练好的值，否则的话，有输入数据，即使不训练，它也会改变权值。这是model中含有BN层和Dropout所带来的的性质。

而在模型测试中，通常还会看见torch.no_grad()，它有什么作用呢？

  • 用于停止autograd模块的工作，起到加速和节省显存的作用（具体行为就是停止gradient计算，从而节省了GPU算力和显存）
  • 不会影响dropout和batchnorm层的行为

model.eval()与torch.no_grad()同时用，可以更加节省cpu算力，加速计算。