torch.nn 下的常用深度学习函数

最新推荐文章于 2026-04-04 12:18:22 发布

原创

最新推荐文章于 2026-04-04 12:18:22 发布 · 1.4k 阅读

标签

#深度学习 #人工智能

1. 层（Layers）

这些函数用于定义神经网络中的各种层，是构建模型的基础模块。

torch.nn.Linear
- 用途：全连接层（也称为线性层）。用于将输入数据从一个维度映射到另一个维度，常用于神经网络的隐藏层和输出层。
- 示例：torch.nn.Linear(in_features=10, out_features=5)，将输入维度为10的数据映射到输出维度为5。

in_features：输入特征的数量，即输入张量的最后一个维度的大小。
out_features：输出特征的数量，即输出张量的最后一个维度的大小。
bias：可选参数，默认为True。当设置为True时，会添加一个偏置项。

import torch
import torch.nn as nn

# 创建一个线性层，输入特征数为 10，输出特征数为 5
linear_layer = nn.Linear(in_features=10, out_features=5)

# 创建一个输入张量，形状为 (batch_size, in_features) = (3, 10)
input_tensor = torch.randn(3, 10)

# 通过线性层进行前向传播
output_tensor = linear_layer(input_tensor)

print("输入张量形状:", input_tensor.shape)
print("输出张量形状:", output_tensor.shape)

输入张量形状: torch.Size([3, 10])
输出张量形状: torch.Size([3, 5])

torch.nn.Conv2d
- 用途：二维卷积层。用于处理二维数据（如图像），通过卷积核提取局部特征，常用于卷积神经网络（CNN）。
- 示例：torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3，stride=1，padding=0)，输入通道数为3，输出通道数为16，卷积核大小为3×3。

- in_channels：输入图像的通道数。
- out_channels：卷积操作产生的通道数。
- kernel_size：卷积核的大小，可以是一个整数或一个包含两个整数的元组。
- stride：卷积操作的步长，默认为1。
- padding：输入的每一边补充0的层数，默认为0。
- dilation：卷积核元素之间的间距，默认为1。
- groups：将输入分成的组数，默认为1。
- bias：可选参数，默认为True。当设置为True时，会添加一个偏置项。
- padding_mode：填充模式，默认为zeros。

import torch
import torch.nn as nn

# 创建一个卷积层，输入通道数为 3，输出通道数为 16，卷积核大小为 3x3
conv_layer = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)

# 创建一个输入张量，形状为 (batch_size, in_channels, height, width) = (2, 3, 32, 32)
input_tensor = torch.randn(2, 3, 32, 32)

# 通过卷积层进行前向传播
output_tensor = conv_layer(input_tensor)

print("输入张量形状:", input_tensor.shape)
print("输出张量形状:", output_tensor.shape)

输入张量形状: torch.Size([2, 3, 32, 32])
输出张量形状: torch.Size([2, 16, 32, 32])