快捷方式

torch.nn

它们是构建图的基本模块

Buffer

一种不应被视为模型参数的 Tensor。

Parameter

一种应被视为模块参数的 Tensor。

UninitializedParameter

一个未初始化的参数。

UninitializedBuffer

一个未初始化的 buffer。

容器

Module

所有神经网络模块的基类。

Sequential

一个序列容器。

ModuleList

以列表形式保存子模块。

ModuleDict

以字典形式保存子模块。

ParameterList

以列表形式保存参数。

ParameterDict

以字典形式保存参数。

模块的全局钩子

register_module_forward_pre_hook

注册一个适用于所有模块的前向预钩子。

register_module_forward_hook

注册一个适用于所有模块的全局前向钩子。

register_module_backward_hook

注册一个适用于所有模块的后向钩子。

register_module_full_backward_pre_hook

注册一个适用于所有模块的完整后向预钩子。

register_module_full_backward_hook

注册一个适用于所有模块的后向钩子。

register_module_buffer_registration_hook

注册一个适用于所有模块的 buffer 注册钩子。

register_module_module_registration_hook

注册一个适用于所有模块的 module 注册钩子。

register_module_parameter_registration_hook

注册一个适用于所有模块的 parameter 注册钩子。

卷积层

nn.Conv1d

对由若干输入平面组成的输入信号应用 1D 卷积。

nn.Conv2d

对由若干输入平面组成的输入信号应用 2D 卷积。

nn.Conv3d

对由若干输入平面组成的输入信号应用 3D 卷积。

nn.ConvTranspose1d

对由若干输入平面组成的输入图像应用 1D 转置卷积算子。

nn.ConvTranspose2d

对由若干输入平面组成的输入图像应用 2D 转置卷积算子。

nn.ConvTranspose3d

对由若干输入平面组成的输入图像应用 3D 转置卷积算子。

nn.LazyConv1d

一个 torch.nn.Conv1d 模块,其 in_channels 参数支持延迟初始化。

nn.LazyConv2d

一个 torch.nn.Conv2d 模块,其 in_channels 参数支持延迟初始化。

nn.LazyConv3d

一个 torch.nn.Conv3d 模块,其 in_channels 参数支持延迟初始化。

nn.LazyConvTranspose1d

一个 torch.nn.ConvTranspose1d 模块,其 in_channels 参数支持延迟初始化。

nn.LazyConvTranspose2d

一个 torch.nn.ConvTranspose2d 模块,其 in_channels 参数支持延迟初始化。

nn.LazyConvTranspose3d

一个 torch.nn.ConvTranspose3d 模块,其 in_channels 参数支持延迟初始化。

nn.Unfold

从批量输入张量中提取滑动的局部块。

nn.Fold

将滑动的局部块数组组合成一个大的包含张量。

池化层

nn.MaxPool1d

对由若干输入平面组成的输入信号应用 1D 最大池化。

nn.MaxPool2d

对由若干输入平面组成的输入信号应用 2D 最大池化。

nn.MaxPool3d

对由若干输入平面组成的输入信号应用 3D 最大池化。

nn.MaxUnpool1d

计算 MaxPool1d 的部分逆运算。

nn.MaxUnpool2d

计算 MaxPool2d 的部分逆运算。

nn.MaxUnpool3d

计算 MaxPool3d 的部分逆运算。

nn.AvgPool1d

对由若干输入平面组成的输入信号应用 1D 平均池化。

nn.AvgPool2d

对由若干输入平面组成的输入信号应用 2D 平均池化。

nn.AvgPool3d

对由若干输入平面组成的输入信号应用 3D 平均池化。

nn.FractionalMaxPool2d

对由若干输入平面组成的输入信号应用 2D 分数最大池化。

nn.FractionalMaxPool3d

对由若干输入平面组成的输入信号应用 3D 分数最大池化。

nn.LPPool1d

对由若干输入平面组成的输入信号应用 1D 幂平均池化。

nn.LPPool2d

对由若干输入平面组成的输入信号应用 2D 幂平均池化。

nn.LPPool3d

对由若干输入平面组成的输入信号应用 3D 幂平均池化。

nn.AdaptiveMaxPool1d

对由若干输入平面组成的输入信号应用 1D 自适应最大池化。

nn.AdaptiveMaxPool2d

对由若干输入平面组成的输入信号应用 2D 自适应最大池化。

nn.AdaptiveMaxPool3d

对由若干输入平面组成的输入信号应用 3D 自适应最大池化。

nn.AdaptiveAvgPool1d

对由若干输入平面组成的输入信号应用 1D 自适应平均池化。

nn.AdaptiveAvgPool2d

对由若干输入平面组成的输入信号应用 2D 自适应平均池化。

nn.AdaptiveAvgPool3d

对由若干输入平面组成的输入信号应用 3D 自适应平均池化。

填充层

nn.ReflectionPad1d

使用输入边界的反射来填充输入张量。

nn.ReflectionPad2d

使用输入边界的反射来填充输入张量。

nn.ReflectionPad3d

使用输入边界的反射来填充输入张量。

nn.ReplicationPad1d

使用输入边界的复制来填充输入张量。

nn.ReplicationPad2d

使用输入边界的复制来填充输入张量。

nn.ReplicationPad3d

使用输入边界的复制来填充输入张量。

nn.ZeroPad1d

使用零填充输入张量的边界。

nn.ZeroPad2d

使用零填充输入张量的边界。

nn.ZeroPad3d

使用零填充输入张量的边界。

nn.ConstantPad1d

使用常量值填充输入张量的边界。

nn.ConstantPad2d

使用常量值填充输入张量的边界。

nn.ConstantPad3d

使用常量值填充输入张量的边界。

nn.CircularPad1d

使用输入边界的循环填充来填充输入张量。

nn.CircularPad2d

使用输入边界的循环填充来填充输入张量。

nn.CircularPad3d

使用输入边界的循环填充来填充输入张量。

非线性激活函数 (加权和相关)

nn.ELU

逐元素应用 Exponential Linear Unit (ELU) 函数。

nn.Hardshrink

逐元素应用 Hard Shrinkage (Hardshrink) 函数。

nn.Hardsigmoid

逐元素应用 Hardsigmoid 函数。

nn.Hardtanh

逐元素应用 HardTanh 函数。

nn.Hardswish

逐元素应用 Hardswish 函数。

nn.LeakyReLU

逐元素应用 LeakyReLU 函数。

nn.LogSigmoid

逐元素应用 Logsigmoid 函数。

nn.MultiheadAttention

允许模型共同关注来自不同表示子空间的信息。

nn.PReLU

逐元素应用 PReLU 函数。

nn.ReLU

逐元素应用整流线性单元函数。

nn.ReLU6

逐元素应用 ReLU6 函数。

nn.RReLU

逐元素应用随机 leaky 整流线性单元函数。

nn.SELU

逐元素应用 SELU 函数。

nn.CELU

逐元素应用 CELU 函数。

nn.GELU

应用高斯误差线性单元函数。

nn.Sigmoid

逐元素应用 Sigmoid 函数。

nn.SiLU

逐元素应用 Sigmoid Linear Unit (SiLU) 函数。

nn.Mish

逐元素应用 Mish 函数。

nn.Softplus

逐元素应用 Softplus 函数。

nn.Softshrink

逐元素应用 soft shrinkage 函数。

nn.Softsign

逐元素应用 Softsign 函数。

nn.Tanh

逐元素应用 Hyperbolic Tangent (Tanh) 函数。

nn.Tanhshrink

逐元素应用 Tanhshrink 函数。

nn.Threshold

对输入 Tensor 的每个元素进行阈值处理。

nn.GLU

应用门控线性单元函数。

非线性激活函数 (其他)

nn.Softmin

对 n 维输入 Tensor 应用 Softmin 函数。

nn.Softmax

对 n 维输入 Tensor 应用 Softmax 函数。

nn.Softmax2d

对每个空间位置的特征应用 SoftMax。

nn.LogSoftmax

对 n 维输入 Tensor 应用 log(Softmax(x))\log(\text{Softmax}(x)) 函数。

nn.AdaptiveLogSoftmaxWithLoss

高效的 softmax 近似。

归一化层

nn.BatchNorm1d

对 2D 或 3D 输入应用批归一化。

nn.BatchNorm2d

对 4D 输入应用批归一化。

nn.BatchNorm3d

对 5D 输入应用批归一化。

nn.LazyBatchNorm1d

一个支持延迟初始化的 torch.nn.BatchNorm1d 模块。

nn.LazyBatchNorm2d

一个支持延迟初始化的 torch.nn.BatchNorm2d 模块。

nn.LazyBatchNorm3d

一个支持延迟初始化的 torch.nn.BatchNorm3d 模块。

nn.GroupNorm

对 mini-batch 输入应用组归一化。

nn.SyncBatchNorm

对 N 维输入应用批归一化。

nn.InstanceNorm1d

应用实例归一化。

nn.InstanceNorm2d

应用实例归一化。

nn.InstanceNorm3d

应用实例归一化。

nn.LazyInstanceNorm1d

一个 torch.nn.InstanceNorm1d 模块,其 num_features 参数支持延迟初始化。

nn.LazyInstanceNorm2d

一个 torch.nn.InstanceNorm2d 模块,其 num_features 参数支持延迟初始化。

nn.LazyInstanceNorm3d

一个 torch.nn.InstanceNorm3d 模块,其 num_features 参数支持延迟初始化。

nn.LayerNorm

对 mini-batch 输入应用层归一化。

nn.LocalResponseNorm

对输入信号应用局部响应归一化。

nn.RMSNorm

对 mini-batch 输入应用均方根层归一化。

循环层

nn.RNNBase

RNN 模块 (RNN, LSTM, GRU) 的基类。

nn.RNN

对输入序列应用带有 tanh\tanhReLU\text{ReLU} 非线性的多层 Elman RNN。

nn.LSTM

对输入序列应用多层长短期记忆 (LSTM) RNN。

nn.GRU

对输入序列应用多层门控循环单元 (GRU) RNN。

nn.RNNCell

一个带有 tanh 或 ReLU 非线性的 Elman RNN cell。

nn.LSTMCell

一个长短期记忆 (LSTM) cell。

nn.GRUCell

一个门控循环单元 (GRU) cell。

Transformer 层

nn.Transformer

一个 transformer 模型。

nn.TransformerEncoder

TransformerEncoder 是 N 个编码器层的堆叠。

nn.TransformerDecoder

TransformerDecoder 是 N 个解码器层的堆叠。

nn.TransformerEncoderLayer

TransformerEncoderLayer 由自注意力和前馈网络组成。

nn.TransformerDecoderLayer

TransformerDecoderLayer 由自注意力、多头注意力和前馈网络组成。

线性层

nn.Identity

一个占位符恒等算子,对参数不敏感。

nn.Linear

对输入数据应用仿射线性变换:y=xAT+by = xA^T + b

nn.Bilinear

对输入数据应用双线性变换:y=x1TAx2+by = x_1^T A x_2 + b

nn.LazyLinear

一个 torch.nn.Linear 模块,其 in_features 参数是推断的。

Dropout 层

nn.Dropout

在训练期间,以概率 p 随机将输入张量的一些元素归零。

nn.Dropout1d

随机将整个通道归零。

nn.Dropout2d

随机将整个通道归零。

nn.Dropout3d

随机将整个通道归零。

nn.AlphaDropout

对输入应用 Alpha Dropout。

nn.FeatureAlphaDropout

随机遮蔽整个通道。

稀疏层

nn.Embedding

一个简单的查找表,用于存储固定词典和大小的嵌入。

nn.EmbeddingBag

计算嵌入“包”的总和或平均值,无需实例化中间嵌入。

距离函数

nn.CosineSimilarity

返回 x1x_1x2x_2 之间的余弦相似度,沿 dim 计算。

nn.PairwiseDistance

计算输入向量之间或输入矩阵列之间的成对距离。

损失函数

nn.L1Loss

创建一个准则,用于衡量输入 xx 和目标 yy 中每个元素之间的平均绝对误差 (MAE)。

nn.MSELoss

创建一个准则,用于衡量输入 xx 和目标 yy 中每个元素之间的均方误差(平方 L2 范数)。

nn.CrossEntropyLoss

此准则计算输入 logit 和目标之间的交叉熵损失。

nn.CTCLoss

连接主义时间分类损失。

nn.NLLLoss

负对数似然损失。

nn.PoissonNLLLoss

目标服从泊松分布的负对数似然损失。

nn.GaussianNLLLoss

高斯负对数似然损失。

nn.KLDivLoss

Kullback-Leibler 散度损失。

nn.BCELoss

创建一个准则,用于衡量目标和输入概率之间的二元交叉熵。

nn.BCEWithLogitsLoss

此损失函数在一个类中结合了 Sigmoid 层和 BCELoss

nn.MarginRankingLoss

创建一个准则,用于衡量给定输入 x1x1、$x2$(两个 1D 小批量或 0D 张量)和标签 yy(一个 1D 小批量或 0D 张量,包含 1 或 -1)时的损失。

nn.HingeEmbeddingLoss

衡量给定输入张量 xx 和标签张量 yy(包含 1 或 -1)时的损失。

nn.MultiLabelMarginLoss

创建一个准则,用于优化输入 xx(一个 2D 小批量 张量)和输出 yy(一个 2D 目标类别索引 张量)之间的多类别多分类铰链损失(基于间隔的损失)。

nn.HuberLoss

创建一个准则,当逐元素的绝对误差小于 delta 时使用平方项,否则使用 delta 放大的 L1 项。

nn.SmoothL1Loss

创建一个准则,当逐元素的绝对误差小于 beta 时使用平方项,否则使用 L1 项。

nn.SoftMarginLoss

创建一个准则,用于优化输入张量 xx 和目标张量 yy(包含 1 或 -1)之间的两类别分类逻辑损失。

nn.MultiLabelSoftMarginLoss

创建一个准则,用于优化基于最大熵的多标签一对多损失,该损失衡量输入 xx 和大小为 (N,C)(N, C) 的目标 yy 之间的差异。

nn.CosineEmbeddingLoss

创建一个准则,用于衡量给定输入张量 x1x_1、$x_2$ 和值为 1 或 -1 的 张量 标签 yy 时的损失。

nn.MultiMarginLoss

创建一个准则,用于优化输入 xx(一个 2D 小批量 张量)和输出 yy(一个 1D 目标类别索引张量,$0 \leq y \leq \text{x.size}(1)-1$)之间的多类别分类铰链损失(基于间隔的损失)。

nn.TripletMarginLoss

创建一个准则,用于衡量给定输入张量 x1x1、$x2$、$x3$ 和一个大于 $0$ 的间隔(margin)时的三元组损失。

nn.TripletMarginWithDistanceLoss

创建一个准则,用于衡量给定输入张量 aa、$p$ 和 nn(分别代表锚点、正例和负例)以及一个非负实值函数(“距离函数”)时的三元组损失,该函数用于计算锚点与正例之间的关系(“正距离”)以及锚点与负例之间的关系(“负距离”)。

视觉层

nn.PixelShuffle

根据上采样因子重新排列张量中的元素。

nn.PixelUnshuffle

反转 PixelShuffle 操作。

nn.Upsample

对给定的多通道 1D(时间)、2D(空间)或 3D(体积)数据进行上采样。

nn.UpsamplingNearest2d

对由多个输入通道组成的输入信号应用 2D 最近邻上采样。

nn.UpsamplingBilinear2d

对由多个输入通道组成的输入信号应用 2D 双线性上采样。

混洗层

nn.ChannelShuffle

划分并重新排列张量中的通道。

数据并行层(多 GPU,分布式)

nn.DataParallel

在模块级别实现数据并行。

nn.parallel.DistributedDataParallel

在模块级别基于 torch.distributed 实现分布式数据并行。

工具函数

来自 torch.nn.utils 模块

用于裁剪参数梯度的工具函数。

clip_grad_norm_

裁剪可迭代参数的梯度范数。

clip_grad_norm

裁剪可迭代参数的梯度范数。

clip_grad_value_

按指定值裁剪可迭代参数的梯度。

get_total_norm

计算可迭代张量的范数。

clip_grads_with_norm_

根据预计算的总范数和期望的最大范数,缩放可迭代参数的梯度。

用于将模块参数展平为单个向量以及从单个向量还原的工具函数。

parameters_to_vector

将可迭代参数展平为单个向量。

vector_to_parameters

将向量的切片复制到可迭代参数中。

用于融合 Module 和 BatchNorm 模块的工具函数。

fuse_conv_bn_eval

将卷积模块和 BatchNorm 模块融合成一个新的卷积模块。

fuse_conv_bn_weights

将卷积模块参数和 BatchNorm 模块参数融合成新的卷积模块参数。

fuse_linear_bn_eval

将线性模块和 BatchNorm 模块融合成一个新的线性模块。

fuse_linear_bn_weights

将线性模块参数和 BatchNorm 模块参数融合成新的线性模块参数。

用于转换模块参数内存格式的工具函数。

convert_conv2d_weight_memory_format

nn.Conv2d.weightmemory_format 转换为指定的 memory_format

convert_conv3d_weight_memory_format

nn.Conv3d.weightmemory_format 转换为指定的 memory_format。此转换递归应用于嵌套的 nn.Module,包括 module

用于对模块参数应用和移除权重归一化的工具函数。

weight_norm

对给定模块中的参数应用权重归一化。

remove_weight_norm

从模块中移除权重归一化重参数化。

spectral_norm

对给定模块中的参数应用谱归一化。

remove_spectral_norm

从模块中移除谱归一化重参数化。

用于初始化模块参数的工具函数。

skip_init

给定模块类对象和 args / kwargs,实例化模块而不初始化参数 / 缓冲区。

用于修剪模块参数的工具类和函数。

prune.BasePruningMethod

创建新修剪技术的抽象基类。

prune.PruningContainer

包含一系列修剪方法用于迭代修剪的容器。

prune.Identity

不修剪任何单元但生成带全 1 掩码的修剪重参数化的实用修剪方法。

prune.RandomUnstructured

随机修剪张量中(当前未修剪的)单元。

prune.L1Unstructured

通过将 L1 范数最低的单元归零来修剪张量中(当前未修剪的)单元。

prune.RandomStructured

随机修剪张量中整个(当前未修剪的)通道。

prune.LnStructured

根据 Ln 范数修剪张量中整个(当前未修剪的)通道。

prune.CustomFromMask

prune.identity

应用修剪重参数化而不修剪任何单元。

prune.random_unstructured

通过移除随机的(当前未修剪的)单元来修剪张量。

prune.l1_unstructured

通过移除 L1 范数最低的单元来修剪张量。

prune.random_structured

沿指定维度移除随机通道来修剪张量。

prune.ln_structured

沿指定维度移除 Ln 范数最低的通道来修剪张量。

prune.global_unstructured

通过应用指定的 pruning_method,全局修剪 parameters 中对应于所有参数的张量。

prune.custom_from_mask

通过应用 mask 中预计算的掩码,修剪 module 中名为 name 的参数对应的张量。

prune.remove

从模块中移除修剪重参数化,并从前向钩子中移除修剪方法。

prune.is_pruned

通过查找修剪前置钩子检查模块是否已修剪。

使用 torch.nn.utils.parameterize.register_parametrization() 中的新参数化功能实现的参数化(Parametrizations)。

parametrizations.orthogonal

对矩阵或一批矩阵应用正交或酉参数化。

parametrizations.weight_norm

对给定模块中的参数应用权重归一化。

parametrizations.spectral_norm

对给定模块中的参数应用谱归一化。

用于对现有模块中的张量进行参数化的工具函数。请注意,这些函数可用于根据一个特定的函数对给定参数或缓冲区进行参数化,该函数将输入空间映射到参数化空间。它们并不是将对象转换为参数的参数化。有关如何实现您自己的参数化的更多信息,请参阅参数化教程

parametrize.register_parametrization

在模块中向张量注册参数化。

parametrize.remove_parametrizations

移除模块中张量上的参数化。

parametrize.cached

上下文管理器,可在使用 register_parametrization() 注册的参数化中启用缓存系统。

parametrize.is_parametrized

确定模块是否具有参数化。

parametrize.ParametrizationList

一个序列容器,用于保存和管理参数化 torch.nn.Module 的原始参数或缓冲区。

用于以无状态方式调用给定模块的工具函数。

stateless.functional_call

通过将模块参数和缓冲区替换为提供的参数和缓冲区,对模块执行函数式调用。

其他模块中的工具函数

nn.utils.rnn.PackedSequence

保存打包序列的数据和 batch_sizes 列表。

nn.utils.rnn.pack_padded_sequence

打包包含可变长度填充序列的张量。

nn.utils.rnn.pad_packed_sequence

填充打包的可变长度序列批次。

nn.utils.rnn.pad_sequence

使用 padding_value 填充可变长度张量列表。

nn.utils.rnn.pack_sequence

打包可变长度张量列表。

nn.utils.rnn.unpack_sequence

将 PackedSequence 解包为可变长度张量列表。

nn.utils.rnn.unpad_sequence

将填充张量去除填充为可变长度张量列表。

nn.Flatten

将连续的维度范围展平为一个张量。

nn.Unflatten

还原张量维度,将其扩展到所需的形状。

量化函数

量化是指以低于浮点精度的位宽执行计算和存储张量的技术。PyTorch 支持逐张量和逐通道的非对称线性量化。要了解如何在 PyTorch 中使用量化函数,请参阅量化文档。

惰性模块初始化

nn.modules.lazy.LazyModuleMixin

用于惰性初始化参数的模块的混合类,也称为“惰性模块”。

别名

以下是 torch.nn 中对应部分的别名。

nn.modules.normalization.RMSNorm

对 mini-batch 输入应用均方根层归一化。

文档

访问 PyTorch 的完整开发者文档

查看文档

教程

获取针对初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源