PyTorch之nn.Module与nn.functional用法区别

1. nn.Module

在PyTorch中，nn.Module 类扮演着核心角色，它是构建任何自定义神经网络层、复杂模块或完整神经网络架构的基础构建块。通过继承 nn.Module 并在其子类中定义模型结构和前向传播逻辑（forward() 方法），开发者能够方便地搭建并训练深度学习模型。

关于 nn.Module 的更多介绍可以参考博客：PyTorch之nn.Module、nn.Sequential、nn.ModuleList使用详解

这里，我们基于nn.Module创建一个简单的神经网络模型，实现代码如下：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(MyModel, self).__init__()
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.layer2 = nn.Linear(hidden_size, output_size)
        
    def forward(self, x):
        x = torch.relu(self.layer1(x))
        x = self.layer2(x)
        return x

2. nn.functional

nn.functional 是PyTorch中一个重要的模块，它包含了许多用于构建神经网络的函数。与 nn.Module 不同，nn.functional 中的函数不具有可学习的参数。这些函数通常用于执行各种非线性操作、损失函数、激活函数等。

2.1 基本用法

如何在神经网络中使用nn.functional？

在PyTorch中，你可以轻松地在神经网络中使用 nn.functional 函数。通常，你只需将输入数据传递给这些函数，并将它们作为网络的一部分。

以下是一个简单的示例，演示如何在一个全连接神经网络中使用ReLU激活函数：

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

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(64, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

在上述示例中，我们首先导入nn.functional 模块，然后在网络的forward 方法中使用F.relu 函数作为激活函数。

nn.functional 的主要优势是它的计算效率和灵活性，因为它允许你以函数的方式直接调用这些操作，而不需要创建额外的层。

2.2 常用函数

（1）激活函数

激活函数是神经网络中的关键组件，它们引入非线性性，使网络能够拟合复杂的数据。以下是一些常见的激活函数：

• ReLU（Rectified Linear Unit）
  ReLU是一种简单而有效的激活函数，它将输入值小于零的部分设为零，大于零的部分保持不变。它的数学表达式如下：

output = F.relu(input)

• Sigmoid
  Sigmoid函数将输入值映射到0和1之间，常用于二分类问题的输出层。它的数学表达式如下：

output = F.sigmoid(input)

• Tanh（双曲正切）
  Tanh函数将输入值映射到-1和1之间，它具有零中心化的特性，通常在循环神经网络中使用。它的数学表达式如下：

output = F.tanh(input)

（2）损失函数

损失函数用于度量模型的预测与真实标签之间的差距。PyTorch的nn.functional 模块包含了各种常用的损失函数，例如：

• 交叉熵损失（Cross-Entropy Loss）
  交叉熵损失通常用于多分类问题，计算模型的预测分布与真实分布之间的差异。它的数学表达式如下：

loss = F.cross_entropy(input, target)

• 均方误差损失（Mean Squared Error Loss）
  均方误差损失通常用于回归问题，度量模型的预测值与真实值之间的平方差。它的数学表达式如下：

loss = F.mse_loss(input, target)

• L1 损失
  L1损失度量预测值与真实值之间的绝对差距，通常用于稀疏性正则化。它的数学表达式如下：

loss = F.l1_loss(input, target)

（3）非线性操作

nn.functional 模块还包含了许多非线性操作，如池化、归一化等。

• 最大池化（Max Pooling）
  最大池化是一种用于减小特征图尺寸的操作，通常用于卷积神经网络中。它的数学表达式如下：

output = F.max_pool2d(input, kernel_size)

• 批量归一化（Batch Normalization）
  批量归一化是一种用于提高训练稳定性和加速收敛的技术。它的数学表达式如下：

output = F.batch_norm(input, mean, std, weight, bias)

3. nn.Module 与 nn.functional

3.1 主要区别

nn.Module 与 nn.functional 的主要区别在于：

• nn.Module实现的layers是一个特殊的类，都是由class Layer(nn.Module)定义，会自动提取可学习的参数；
• nn.functional中的函数更像是纯函数，由def function(input)定义。

注意：

1. 如果模型有可学习的参数时，最好使用nn.Module。
2. 激活函数（ReLU、sigmoid、Tanh)、池化(MaxPool)等层没有可学习的参数,可以使用对应的functional函数。
3. 卷积、全连接等有可学习参数的网络建议使用nn.Module。
4. dropout没有可学习参数，但建议使用nn.Dropout而不是nn.functional.dropout。

3.2 具体样例：nn.ReLU() 与 F.relu()

nn.ReLU() ：

import torch.nn as nn
'''
nn.ReLU()

F.relu()：

import torch.nn.functional as F
'''
out = F.relu(input)

其实这两种方法都是使用relu激活，只是使用的场景不一样，F.relu()是函数调用，一般使用在foreward函数里。而nn.ReLU()是模块调用，一般在定义网络层的时候使用。

当用print(net)输出时，nn.ReLU()会有对应的层，而F.ReLU()是没有输出的。

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

class NET1(nn.Module):
    def __init__(self):
        super(NET1, self).__init__()
        self.conv = nn.Conv2d(3, 16, 3, 1, 1)
        self.bn = nn.BatchNorm2d(16)
        self.relu = nn.ReLU()  # 模块的激活函数

    def forward(self, x):
        out = self.conv(x)
        x = self.bn(x)
        out = self.relu()
        return out


class NET2(nn.Module):
    def __init__(self):
        super(NET2, self).__init__()
        self.conv = nn.Conv2d(3, 16, 3, 1, 1)
        self.bn = nn.BatchNorm2d(16)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        out = F.relu(x)  # 函数的激活函数
        return out


net1 = NET1()
net2 = NET2()
print(net1)
print(net2)

参考资料

• PyTorch的nn.Module类的详细介绍
• PyTorch nn.functional 模块详解：探索神经网络的魔法工具箱
• pytorch：F.relu() 与 nn.ReLU() 的区别