torch.nn.functional的用法

介绍

torch.nn.functional 是 PyTorch 中的一个模块，提供了许多函数式的神经网络操作，包括激活函数、损失函数、卷积操作等。这些函数是无状态的（stateless），与 torch.nn 中的模块化层（如 nn.ReLU、nn.Conv2d 等）不同，torch.nn.functional 提供的是直接的函数调用方式。

激活函数

torch.nn.functional 提供了许多常用的激活函数，例如 ReLU、Sigmoid、Tanh 等。

import torch.nn.functional as F

示例

import torch  
import torch.nn.functional as F  

x = torch.tensor([-1.0, 0.0, 1.0])  
relu_output = F.relu(x)  # ReLU 激活  
softmax_output = F.softmax(x, dim=0)  # Softmax 激活  
print(relu_output)  # tensor([0., 0., 1.])  
print(softmax_output)  # tensor([0.0900, 0.2447, 0.6652])

损失函数

torch.nn.functional 提供了许多损失函数，与 torch.nn 中的模块化损失函数（如 nn.CrossEntropyLoss）功能相同，但需要显式传入参数。

示例

input = torch.tensor([[0.5, 1.5], [2.0, 1.0]], requires_grad=True)  
target = torch.tensor([1, 0])  
loss = F.cross_entropy(input, target)  # 交叉熵损失  
print(loss)  # tensor(1.2412, grad_fn=<NllLossBackward>)

卷积操作

torch.nn.functional 提供了卷积操作的函数式实现，例如 F.conv1d、F.conv2d、F.conv3d。

示例

input = torch.randn(1, 1, 5)  # 输入：batch_size=1, channels=1, width=5  
weight = torch.randn(1, 1, 3)  # 卷积核：out_channels=1, in_channels=1, kernel_size=3  
output = F.conv1d(input, weight)  
print(output.shape)  # torch.Size([1, 1, 3])

池化

torch.nn.functional 提供了池化操作的函数式实现，例如最大池化和平均池化。

示例

input = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]])  # 输入：batch_size=1, channels=1, height=2, width=2  
output = F.max_pool2d(input, kernel_size=2)  
print(output)  # tensor([[[[4.]]]])

归一化操作

torch.nn.functional 提供了归一化操作的函数式实现，例如 BatchNorm、LayerNorm 等。

示例

input = torch.randn(2, 3)  # 输入：batch_size=2, features=3  
output = F.layer_norm(input, normalized_shape=(3,))  
print(output)

Dropout

torch.nn.functional 提供了 Dropout 的函数式实现。

示例

input = torch.tensor([1.0, 2.0, 3.0])  
output = F.dropout(input, p=0.5, training=True)  # 50% 概率随机置零  
print(output)

torch.nn.functional 与 torch.nn 的区别