torch.nn.functional

torch.nn.functional.conv1d

是

torch.nn.functional.conv2d

是

torch.nn.functional.conv3d

否

torch.nn.functional.conv_transpose1d

是

只支持float16

torch.nn.functional.conv_transpose2d

是

torch.nn.functional.conv_transpose3d

否

torch.nn.functional.unfold

是

torch.nn.functional.fold

是

torch.nn.functional.avg_pool1d

否

torch.nn.functional.avg_pool2d

否

torch.nn.functional.avg_pool3d

否

torch.nn.functional.max_pool1d

否

torch.nn.functional.max_pool2d

是

只支持dialtion=1的场景

torch.nn.functional.max_pool3d

否

torch.nn.functional.max_unpool1d

是

torch.nn.functional.max_unpool2d

是

torch.nn.functional.max_unpool3d

是

torch.nn.functional.lp_pool1d

否

torch.nn.functional.lp_pool2d

否

torch.nn.functional.adaptive_max_pool1d

torch.nn.functional.adaptive_max_pool2d

是

输入shape要能够整除输出shape

torch.nn.functional.adaptive_max_pool3d

否

torch.nn.functional.adaptive_avg_pool1d

是

只支持float16，float32

torch.nn.functional.adaptive_avg_pool2d

是

只支持float16，float32

torch.nn.functional.adaptive_avg_pool3d

是

仅支持D=1 && H=1 && W=1

torch.nn.functional.fractional_max_pool2d

否

torch.nn.functional.fractional_max_pool3d

否

torch.nn.functional.scaled_dot_product_attention

否

torch.nn.functional.threshold

是

torch.nn.functional.threshold_

是

torch.nn.functional.relu

是

torch.nn.functional.relu_

是

torch.nn.functional.hardtanh

是

torch.nn.functional.hardtanh_

是

torch.nn.functional.hardswish

是

torch.nn.functional.relu6

是

torch.nn.functional.elu

是

torch.nn.functional.elu_

是

torch.nn.functional.selu

是

torch.nn.functional.celu

是

torch.nn.functional.leaky_relu

是

torch.nn.functional.leaky_relu_

是

torch.nn.functional.prelu

是

torch.nn.functional.rrelu

是

torch.nn.functional.rrelu_

是

torch.nn.functional.glu

否

torch.nn.functional.gelu

否

torch.nn.functional.logsigmoid

是

torch.nn.functional.hardshrink

是

torch.nn.functional.tanhshrink

是

torch.nn.functional.softsign

是

torch.nn.functional.softplus

是

torch.nn.functional.softmin

是

torch.nn.functional.softmax

是

不支持uint、int8、int16、int32、int64、bool

torch.nn.functional.softshrink

是

torch.nn.functional.gumbel_softmax

torch.nn.functional.log_softmax

是

不支持uint、int8、int16、int32、int64、bool

torch.nn.functional.tanh

是

torch.nn.functional.sigmoid

是

torch.nn.functional.hardsigmoid

是

torch.nn.functional.silu

是

torch.nn.functional.mish

否

torch.nn.functional.batch_norm

是

torch.nn.functional.group_norm

否

torch.nn.functional.instance_norm

是

torch.nn.functional.layer_norm

是

torch.nn.functional.local_response_norm

否

torch.nn.functional.normalize

否

torch.nn.functional.linear

否

torch.nn.functional.bilinear

是

Atlas A2 训练系列产品：不支持int32输入类型

torch.nn.functional.dropout

是

torch.nn.functional.alpha_dropout

是

torch.nn.functional.feature_alpha_dropout

是

torch.nn.functional.dropout1d

否

torch.nn.functional.dropout2d

是

torch.nn.functional.dropout3d

是

torch.nn.functional.embedding

是

torch.nn.functional.embedding_bag

是

torch.nn.functional.one_hot

是

torch.nn.functional.pairwise_distance

是

torch.nn.functional.cosine_similarity

是

torch.nn.functional.pdist

是

torch.nn.functional.binary_cross_entropy

是

torch.nn.functional.binary_cross_entropy_with_logits

是

torch.nn.functional.poisson_nll_loss

是

不支持int64

torch.nn.functional.cosine_embedding_loss

是

不支持uint8,int8,int16,int32,int64,bool

torch.nn.functional.cross_entropy

是

torch.nn.functional.ctc_loss

是

不支持float64

torch.nn.functional.gaussian_nll_loss

是

torch.nn.functional.hinge_embedding_loss

是

torch.nn.functional.kl_div

否

torch.nn.functional.l1_loss

是

torch.nn.functional.mse_loss

是

torch.nn.functional.margin_ranking_loss

是

torch.nn.functional.multilabel_margin_loss

是

torch.nn.functional.multilabel_soft_margin_loss

否

torch.nn.functional.multi_margin_loss

否

可以走CPU实现

torch.nn.functional.nll_loss

是

torch.nn.functional.huber_loss

否

torch.nn.functional.smooth_l1_loss

是

torch.nn.functional.soft_margin_loss

否

torch.nn.functional.triplet_margin_loss

否

torch.nn.functional.triplet_margin_with_distance_loss

否

torch.nn.functional.pixel_shuffle

是

不支持8D输入

torch.nn.functional.pixel_unshuffle

是

torch.nn.functional.pad

是

在输入x为六维以上时可能会出现性能下降问题。

torch.nn.functional.interpolate

是

不支持recompute_scale_factor=True。

torch.nn.functional.upsample

Atlas 训练系列产品：是

Atlas A2 训练系列产品：否

Atlas 训练系列产品：

不能输入较大shape（数据量百万级别）

在torch.nn.Upsample(size=size, mode='bicubic',align_corners=False)中的size为(1,*)时，torch.nn.functional.upsample中的align_corners参数需要设置为True

建议使用float32类型，否则会影响计算性能

align_corners参数为False时，输入和输出的向量shape需相同，否侧输出结果有误

torch.nn.functional.upsample_nearest

是

torch.nn.functional.upsample_bilinear

是

torch.nn.functional.grid_sample

否

torch.nn.functional.affine_grid

是

torch.nn.parallel.data_parallel

否