评价此页

torch.nn.functional#

创建日期:2019年6月11日 | 最后更新日期:2025年12月8日

卷积函数#

conv1d

对由多个输入平面组成的输入信号进行1D卷积操作。

conv2d

对由多个输入平面组成的输入图像应用 2D 卷积。

conv3d

对由多个输入平面组成的输入图像应用 3D 卷积。

conv_transpose1d

对由多个输入平面组成的输入信号应用 1D 转置卷积算子,有时也称为“去卷积”。

conv_transpose2d

对由多个输入平面组成的输入图像应用 2D 转置卷积算子,有时也称为“去卷积”。

conv_transpose3d

对由多个输入平面组成的输入图像应用 3D 转置卷积算子,有时也称为“去卷积”。

unfold

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

fold

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

池化函数#

avg_pool1d

对由多个输入平面组成的输入信号进行1D平均池化操作。

avg_pool2d

kH×kWkH \times kW 区域内,以 sH×sWsH \times sW 为步长应用 2D 平均池化操作。

avg_pool3d

kT×kH×kWkT \times kH \times kW 区域内,以 sT×sH×sWsT \times sH \times sW 为步长应用 3D 平均池化操作。

max_pool1d

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

max_pool2d

在由多个输入平面组成的输入信号上应用 2D 最大池化。

max_pool3d

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

max_unpool1d

计算 MaxPool1d 的部分逆运算。

max_unpool2d

计算 MaxPool2d 的部分逆运算。

max_unpool3d

计算 MaxPool3d 的部分逆运算。

lp_pool1d

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

lp_pool2d

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

lp_pool3d

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

adaptive_max_pool1d

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

adaptive_max_pool2d

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

adaptive_max_pool3d

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

adaptive_avg_pool1d

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

adaptive_avg_pool2d

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

adaptive_avg_pool3d

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

fractional_max_pool2d

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

fractional_max_pool3d

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

注意力机制#

torch.nn.attention.bias 模块包含专为配合 scaled_dot_product_attention 使用而设计的注意力偏置(attention_biases)。

scaled_dot_product_attention

缩放点积注意力 (query, key, value, attn_mask=None, dropout_p=0.0,

非线性激活函数#

threshold

对输入张量的每个元素应用阈值。

threshold_

threshold() 的原地操作版本。

relu

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

relu_

relu() 的原地操作版本。

hardtanh

逐元素应用 HardTanh 函数。

hardtanh_

hardtanh() 的原地操作版本。

hardswish

逐元素应用 hardswish 函数。

relu6

应用逐元素函数 ReLU6(x)=min(max(0,x),6)\text{ReLU6}(x) = \min(\max(0,x), 6)

elu

逐元素应用指数线性单元(ELU)函数。

elu_

elu() 的原地操作版本。

selu

逐元素应用 SELU(x)=scale(max(0,x)+min(0,α(exp(x)1)))\text{SELU}(x) = scale * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1))),其中 α=1.6732632423543772848170429916717\alpha=1.6732632423543772848170429916717scale=1.0507009873554804934193349852946scale=1.0507009873554804934193349852946

celu

逐元素应用 CELU(x)=max(0,x)+min(0,α(exp(x/α)1))\text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))

leaky_relu

逐元素应用 LeakyReLU(x)=max(0,x)+negative_slopemin(0,x)\text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)

leaky_relu_

leaky_relu() 的原地操作版本。

prelu

逐元素应用函数 PReLU(x)=max(0,x)+weightmin(0,x)\text{PReLU}(x) = \max(0,x) + \text{weight} * \min(0,x),其中 weight 是一个可学习参数。

rrelu

随机化 leaky ReLU。

rrelu_

rrelu() 的原地操作版本。

glu

门控线性单元(Gated Linear Unit)。

gelu

当 approximate 参数为 'none' 时,它逐元素应用函数 GELU(x)=xΦ(x)\text{GELU}(x) = x * \Phi(x)

logsigmoid

逐元素应用 LogSigmoid(xi)=log(11+exp(xi))\text{LogSigmoid}(x_i) = \log \left(\frac{1}{1 + \exp(-x_i)}\right)

hardshrink

逐元素应用硬收缩函数

tanhshrink

逐元素应用 Tanhshrink(x)=xTanh(x)\text{Tanhshrink}(x) = x - \text{Tanh}(x)

softsign

逐元素应用函数 SoftSign(x)=x1+x\text{SoftSign}(x) = \frac{x}{1 + |x|}

softplus

逐元素应用函数 Softplus(x)=1βlog(1+exp(βx))\text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x))

softmin

应用 softmin 函数。

softmax

应用 softmax 函数。

softshrink

逐元素应用软收缩函数

gumbel_softmax

从 Gumbel-Softmax 分布中采样(链接 1 链接 2)并可选择进行离散化。

log_softmax

应用 softmax,随后取对数。

tanh

逐元素应用 Tanh(x)=tanh(x)=exp(x)exp(x)exp(x)+exp(x)\text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)}{\exp(x) + \exp(-x)}

sigmoid

应用逐元素函数 Sigmoid(x)=11+exp(x)\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)}

hardsigmoid

逐元素应用 Hardsigmoid 函数。

silu

逐元素应用 Sigmoid 线性单元(SiLU)函数。

mish

逐元素应用 Mish 函数。

batch_norm

对一批数据中的每个通道应用批归一化。

group_norm

对最后若干个维度应用组归一化。

instance_norm

对批次中每个数据样本的每个通道独立应用实例归一化。

layer_norm

对最后若干个维度应用层归一化。

local_response_norm

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

rms_norm

应用均方根层归一化。

normalize

在指定维度上对输入执行 LpL_p 归一化。

线性函数#

linear

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

bilinear

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

Dropout(随机失活)函数#

dropout

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

alpha_dropout

对输入应用 Alpha Dropout。

feature_alpha_dropout

随机屏蔽整个通道(通道即特征图)。

dropout1d

随机将整个通道置零(通道即一维特征图)。

dropout2d

随机将整个通道置零(通道即二维特征图)。

dropout3d

随机将整个通道置零(通道即三维特征图)。

稀疏函数#

embedding

生成一个简单的查找表,用于在固定字典和尺寸中查找嵌入。

embedding_bag

计算嵌入 包(bags) 的总和、平均值或最大值。

one_hot

接收形状为 (*) 索引值的 LongTensor,并返回一个形状为 (*, num_classes) 的张量,其中除了与输入张量最后一个维度值匹配的索引位置为 1 外,其余位置均为 0。

距离函数#

pairwise_distance

详情请参阅 torch.nn.PairwiseDistance

cosine_similarity

返回 x1x2 之间的余弦相似度,沿指定维度计算。

pdist

计算输入中每一对行向量之间的 p-范数距离。

损失函数#

binary_cross_entropy

计算目标值与输入概率之间的二元交叉熵。

binary_cross_entropy_with_logits

计算目标值与输入 Logits 之间的二元交叉熵。

poisson_nll_loss

计算泊松负对数似然损失。

cosine_embedding_loss

计算余弦嵌入损失。

cross_entropy

计算输入 Logits 与目标值之间的交叉熵损失。

ctc_loss

计算联结时序分类(Connectionist Temporal Classification)损失。

gaussian_nll_loss

计算高斯负对数似然损失。

hinge_embedding_loss

计算 Hinge 嵌入损失。

kl_div

计算 KL 散度损失。

l1_loss

计算 L1 损失(可选加权)。

mse_loss

计算元素级均方误差(可选加权)。

margin_ranking_loss

计算间隔排序损失。

multilabel_margin_loss

计算多标签间隔损失。

multilabel_soft_margin_loss

计算多标签软间隔损失。

multi_margin_loss

计算多重间隔损失(可选加权)。

nll_loss

计算负对数似然损失。

huber_loss

计算 Huber 损失(可选加权)。

smooth_l1_loss

计算平滑 L1 损失。

soft_margin_loss

计算软间隔损失。

triplet_margin_loss

计算给定输入张量之间的三元组损失,且间隔大于 0。

triplet_margin_with_distance_loss

使用自定义距离函数计算输入张量的三元组间隔损失。

视觉函数#

pixel_shuffle

将形状为 (,C×r2,H,W)(*, C \times r^2, H, W) 的张量元素重排为形状为 (,C,H×r,W×r)(*, C, H \times r, W \times r) 的张量,其中 r 为 upscale_factor

pixel_unshuffle

通过将形状为 (,C,H×r,W×r)(*, C, H \times r, W \times r) 的张量元素重排,来执行与 PixelShuffle 相反的操作,得到形状为 (,C×r2,H,W)(*, C \times r^2, H, W) 的张量,其中 r 为 downscale_factor

pad

填充张量。

interpolate

对输入进行下采样/上采样。

upsample

上采样输入。

upsample_nearest

使用最近邻像素值对输入进行上采样。

upsample_bilinear

使用双线性上采样对输入进行上采样。

grid_sample

计算网格采样(grid sample)。

affine_grid

给定一批仿射矩阵 theta,生成 2D 或 3D 流场(采样网格)。

DataParallel 函数(多 GPU、分布式)#

data_parallel#

torch.nn.parallel.data_parallel

在 device_ids 指定的多个 GPU 上并行评估 module(input)。

低精度函数#

ScalingType

_ScalingType 的别名

SwizzleType

_SwizzleType 的别名

grouped_mm

计算一种分组矩阵乘法,它在各个专家之间共享权重形状,但允许每个专家有不规则的 token 数量,这在混合专家模型(MoE)层中很常见。

scaled_mm

scaled_mm(mat_a, mat_b, scale_a, scale_recipe_a, scale_b, scale_recipe_b, swizzle_a, swizzle_b, bias, output_dtype,

scaled_grouped_mm

scaled_grouped_mm(mat_a, mat_b, scale_a, scale_recipe_a, scale_b, scale_recipe_b, swizzle_a, swizzle_b, bias, offs,