torch.nn

若API未标明支持情况,则代表该API的支持情况待验证。

API名称

是否支持

限制与说明

torch.nn.parameter.Parameter

是

支持fp32

torch.nn.parameter.UninitializedParameter

是

  

torch.nn.parameter.UninitializedBuffer

是

  

torch.nn.Module

是

支持fp32

torch.nn.Module.add_module

是

支持fp32

torch.nn.Module.apply

是

支持fp32

torch.nn.Module.bfloat16

是

支持bf16,fp32,int64

torch.nn.Module.buffers

是

  

torch.nn.Module.children

是

支持fp32

torch.nn.Module.cpu

是

支持fp32

torch.nn.Module.cuda

是

支持fp32

torch.nn.Module.double

是

  

torch.nn.Module.eval

是

支持fp32,int64

torch.nn.Module.extra_repr

是

支持fp32

torch.nn.Module.float

是

支持fp16,fp32,int64

torch.nn.Module.forward

是

支持fp32

torch.nn.Module.get_buffer

是

  

torch.nn.Module.get_extra_state

是

  

torch.nn.Module.get_parameter

是

支持fp32

torch.nn.Module.get_submodule

是

支持fp32

torch.nn.Module.half

是

支持fp16,fp32,int64

torch.nn.Module.load_state_dict

是

支持fp32

torch.nn.Module.modules

是

支持fp32

torch.nn.Module.named_buffers

是

  

torch.nn.Module.named_children

是

支持fp32

torch.nn.Module.named_modules

是

支持fp32

torch.nn.Module.named_parameters

是

  

torch.nn.Module.parameters

是

  

torch.nn.Module.register_backward_hook

是

支持fp32

torch.nn.Module.register_buffer

是

支持fp32

torch.nn.Module.register_forward_hook

是

支持fp32

torch.nn.Module.register_forward_pre_hook

是

支持fp32

torch.nn.Module.register_full_backward_hook

是

支持fp32

torch.nn.Module.register_module

是

支持fp32

torch.nn.Module.register_parameter

是

  

torch.nn.Module.requires_grad_

是

  

torch.nn.Module.set_extra_state

是

  

torch.nn.Module.share_memory

否

  

torch.nn.Module.state_dict

是

支持fp32

torch.nn.Module.to(*args, **kwargs)

是

支持fp32,int64

torch.nn.Module.to(device=None, dtype=None, non_blocking=False)

是

支持fp32,int64

torch.nn.Module.to(dtype, non_blocking=False)

是

支持fp32,int64

torch.nn.Module.to(tensor, non_blocking=False)

是

支持fp32,int64

torch.nn.Module.to(memory_format=torch.channels_last)

否

  

torch.nn.Module.to_empty

是

支持fp32

torch.nn.Module.train

是

支持fp32

torch.nn.Module.type

是

支持fp16,fp32,int64

torch.nn.Module.xpu

否

  

torch.nn.Module.zero_grad

是

支持fp32

torch.nn.Sequential

是

支持fp32

torch.nn.Sequential.append

是

支持fp32

torch.nn.ModuleList

是

支持fp32

torch.nn.ModuleList.append

是

支持fp32

torch.nn.ModuleList.extend

是

支持fp32

torch.nn.ModuleList.insert

是

支持fp32

torch.nn.ModuleDict

是

支持fp32

torch.nn.ModuleDict.clear

是

支持fp32

torch.nn.ModuleDict.items

是

支持fp32

torch.nn.ModuleDict.keys

是

支持fp32

torch.nn.ModuleDict.pop

是

支持fp32

torch.nn.ModuleDict.update

是

支持fp32

torch.nn.ModuleDict.values

是

支持fp32

torch.nn.ParameterList

是

支持fp32

torch.nn.ParameterList.append

是

支持fp32

torch.nn.ParameterList.extend

是

支持fp32

torch.nn.ParameterDict

是

支持fp32

torch.nn.ParameterDict.clear

是

支持fp32

torch.nn.ParameterDict.copy

是

支持fp32

torch.nn.ParameterDict.fromkeys

是

支持fp32

torch.nn.ParameterDict.get

是

支持fp32

torch.nn.ParameterDict.items

是

支持fp32

torch.nn.ParameterDict.keys

是

支持fp32

torch.nn.ParameterDict.pop

是

支持fp32

torch.nn.ParameterDict.popitem

是

支持fp32

torch.nn.ParameterDict.setdefault

是

支持fp32

torch.nn.ParameterDict.update

是

支持fp32

torch.nn.ParameterDict.values

是

支持fp32

torch.nn.modules.module.register_module_forward_pre_hook

是

支持fp32

torch.nn.modules.module.register_module_forward_hook

是

支持fp32

torch.nn.modules.module.register_module_backward_hook

是

支持fp32

torch.nn.modules.module.register_module_full_backward_hook

是

支持fp32

torch.nn.Conv1d

是

支持fp32,fp16

torch.nn.Conv2d

是

支持fp16,fp32

torch.nn.Conv3d

     

torch.nn.ConvTranspose1d

是

支持fp32

torch.nn.ConvTranspose2d

是

支持fp16,fp32

torch.nn.ConvTranspose3d

     

torch.nn.LazyConv1d

是

支持fp32,fp16

torch.nn.LazyConv2d

是

支持fp32,fp16

torch.nn.LazyConv3d

否

  

torch.nn.LazyConvTranspose1d

是

支持fp16

torch.nn.LazyConvTranspose2d

是

支持fp16,fp32

torch.nn.LazyConvTranspose3d

     

torch.nn.Unfold

     

torch.nn.Fold

     

torch.nn.MaxPool1d

     

torch.nn.MaxPool2d

     

torch.nn.MaxPool3d

     

torch.nn.MaxUnpool1d

是

支持fp16,fp32

torch.nn.MaxUnpool2d

是

支持fp16,fp32

torch.nn.MaxUnpool3d

     

torch.nn.AvgPool1d

是

支持fp16,fp32

torch.nn.AvgPool2d

是

支持fp32,fp16

torch.nn.AvgPool3d

     

torch.nn.FractionalMaxPool2d

否

  

torch.nn.FractionalMaxPool3d

否

  

torch.nn.LPPool1d

是

支持fp32,fp16,uint8,int8,int16,int32,int64,bool

torch.nn.LPPool2d

是

支持fp32,fp16,int16,int32,int64,bool

torch.nn.AdaptiveMaxPool1d

     

torch.nn.AdaptiveMaxPool2d

     

torch.nn.AdaptiveMaxPool3d

是

支持fp32,fp64

torch.nn.AdaptiveAvgPool1d

是

支持fp32,fp16

torch.nn.AdaptiveAvgPool2d

是

支持fp32,fp16

torch.nn.AdaptiveAvgPool3d

     

torch.nn.ReflectionPad1d

是

支持fp32,fp16

torch.nn.ReflectionPad2d

是

支持fp32,fp16

torch.nn.ReflectionPad3d

否

  

torch.nn.ReplicationPad1d

是

支持fp32,fp16,complex64,complex128

torch.nn.ReplicationPad2d

是

支持fp32,fp16,complex64,complex128

torch.nn.ReplicationPad3d

否

可以走CPU实现。

torch.nn.ZeroPad2d

是

支持bf16,fp32,fp16,uint8,int8,int16,int32,int64,bool,complex64,complex128

torch.nn.ConstantPad1d

是

支持int8,bool

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

torch.nn.ConstantPad2d

是

支持fp32,fp16,uint8,int8,int16,int32,int64,complex64,complex128

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

torch.nn.ConstantPad3d

是

支持fp32,fp16,uint8,int8,int16,int32,int64,bool,complex64,complex128

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

torch.nn.ELU

是

支持fp32,fp16

torch.nn.Hardshrink

是

支持fp32,fp16

torch.nn.Hardsigmoid

是

支持fp16,fp32,int32

torch.nn.Hardtanh

是

支持fp32,fp16,int8,int16,int32,int64

torch.nn.Hardswish

是

支持fp16,fp32

torch.nn.LeakyReLU

是

支持fp16,fp32,fp64

torch.nn.LogSigmoid

是

支持fp32,fp16

torch.nn.MultiheadAttention

     

torch.nn.MultiheadAttention.forward

     

torch.nn.PReLU

是

支持fp32

torch.nn.ReLU

是

支持bf16,fp16,fp32,int8,uint8,int32,int64

torch.nn.ReLU6

是

支持fp32,fp16,uint8,int8,int16,int32,int64

torch.nn.RReLU

     

torch.nn.SELU

是

支持fp16,fp32,int8,int32,fp64,uint8,int16,int64,bool

torch.nn.CELU

是

支持fp32,fp16

torch.nn.GELU

是

支持bf16,fp32,fp16

torch.nn.Sigmoid

是

支持bf16,fp32,fp16,uint8,int8,int16,int32,int64,bool,complex64,complex128

torch.nn.SiLU

是

支持bf16,fp16,fp32

torch.nn.Mish

是

支持fp32,fp16

torch.nn.Softplus

是

支持fp16,fp32

torch.nn.Softshrink

是

支持bf16,fp16,fp32

torch.nn.Softsign

是

支持bf16,fp32,fp16,uint8,int8,int16,int32,int64

torch.nn.Tanh

是

支持bf16,fp16,fp32,bool

torch.nn.Tanhshrink

是

支持fp32,fp16,uint8,int8,int16,int32,int64

torch.nn.Threshold

是

支持fp16,fp32,int8,uint8,int16,int32,int64

torch.nn.GLU

是

支持fp32,fp16

torch.nn.Softmin

是

支持bf16,fp32,fp16

torch.nn.Softmax

是

支持bf16,fp32,fp16

torch.nn.Softmax2d

是

支持bf16,fp32,fp16

torch.nn.LogSoftmax

是

支持bf16,fp32,fp16

torch.nn.AdaptiveLogSoftmaxWithLoss

否

  

torch.nn.AdaptiveLogSoftmaxWithLoss.log_prob

否

  

torch.nn.AdaptiveLogSoftmaxWithLoss.predict

否

  

torch.nn.BatchNorm1d

是

支持fp32,fp16

torch.nn.BatchNorm2d

是

支持fp32,fp16

torch.nn.BatchNorm3d

是

支持fp32,fp16

torch.nn.LazyBatchNorm1d

     

torch.nn.LazyBatchNorm2d

     

torch.nn.LazyBatchNorm3d

     

torch.nn.GroupNorm

是

支持fp32

torch.nn.SyncBatchNorm

是

支持fp16,fp32

torch.nn.SyncBatchNorm.convert_sync_batchnorm

是

  

torch.nn.InstanceNorm1d

     

torch.nn.InstanceNorm2d

     

torch.nn.InstanceNorm3d

     

torch.nn.LazyInstanceNorm1d

     

torch.nn.LazyInstanceNorm2d

     

torch.nn.LazyInstanceNorm3d

     

torch.nn.LayerNorm

是

支持fp32,fp16

torch.nn.LocalResponseNorm

     

torch.nn.RNNBase

     

torch.nn.RNNBase.flatten_parameters

     

torch.nn.RNN

     

torch.nn.LSTM

是

支持fp32

不支持proj_size参数

torch.nn.GRU

     

torch.nn.RNNCell

     

torch.nn.LSTMCell

     

torch.nn.GRUCell

是

支持fp32,fp16

torch.nn.Transformer

是

支持fp32,fp16

torch.nn.Transformer.forward

     

torch.nn.Transformer.generate_square_subsequent_mask

     

torch.nn.TransformerEncoder

     

torch.nn.TransformerEncoder.forward

     

torch.nn.TransformerDecoder

     

torch.nn.TransformerDecoder.forward

     

torch.nn.TransformerEncoderLayer

     

torch.nn.TransformerEncoderLayer.forward

     

torch.nn.TransformerDecoderLayer

     

torch.nn.TransformerDecoderLayer.forward

     

torch.nn.Identity

是

支持fp32

torch.nn.Linear

是

支持fp32,fp16

torch.nn.Bilinear

是

支持fp32,fp16

torch.nn.LazyLinear

是

支持fp16,fp32

torch.nn.Dropout

     

torch.nn.Dropout2d

是

支持fp16,fp32,int64,bool

torch.nn.Dropout3d

     

torch.nn.AlphaDropout

是

支持fp32,fp16

torch.nn.FeatureAlphaDropout

是

支持fp32,fp16

torch.nn.Embedding

是

支持int32,int64

属性max_norm仅支持非负值

torch.nn.Embedding.from_pretrained

是

支持fp64

torch.nn.EmbeddingBag

是

支持int32,int64

torch.nn.EmbeddingBag.forward

是

支持int64

torch.nn.EmbeddingBag.from_pretrained

是

支持int64

torch.nn.CosineSimilarity

     

torch.nn.PairwiseDistance

     

torch.nn.L1Loss

是

支持fp32,fp16,int64,

torch.nn.MSELoss

是

支持fp32,fp16

torch.nn.CrossEntropyLoss

是

支持fp32,fp16

torch.nn.CTCLoss

是

支持fp32

torch.nn.NLLLoss

是

支持fp32,fp16

torch.nn.PoissonNLLLoss

是

支持bf16,fp32,fp16,int64

torch.nn.GaussianNLLLoss

是

支持bf16,fp32,fp16,int16,int32,int64

torch.nn.KLDivLoss

是

支持fp16,fp32,fp64

torch.nn.BCELoss

是

支持fp32,fp16

torch.nn.BCEWithLogitsLoss

是

支持fp32,fp16,int16,int32,int64,bool

torch.nn.MarginRankingLoss

是

支持bf16,fp32,fp16,int8,int16,int32,int64

torch.nn.HingeEmbeddingLoss

是

支持bf16,fp32,fp16,uint8,int8,int16,int32,int64

torch.nn.MultiLabelMarginLoss

     

torch.nn.HuberLoss

否

  

torch.nn.SmoothL1Loss

是

支持fp32,fp16

torch.nn.SoftMarginLoss

     

torch.nn.MultiLabelSoftMarginLoss

是

支持fp16,fp32

torch.nn.CosineEmbeddingLoss

     

torch.nn.MultiMarginLoss

否

  

torch.nn.TripletMarginLoss

是

支持fp16,fp32,int8,uint8,int16,int32,int64

torch.nn.TripletMarginWithDistanceLoss

是

支持bf16,fp32,fp16

torch.nn.PixelShuffle

是

支持fp16,fp32,fp64,int8,uint8,int16,int32,int64,bool

torch.nn.PixelUnshuffle

是

支持fp16,fp32,fp64,int8,uint8,int16,int32,int64,bool

torch.nn.Upsample

是

支持fp16,fp32,fp64

只支持mode = nearest

torch.nn.UpsamplingNearest2d

是

支持fp32,fp16,uint8

torch.nn.UpsamplingBilinear2d

     

torch.nn.ChannelShuffle

是

支持bf16,fp32,fp16,uint8,int8,int16,int32,int64,bool,complex64,complex128

torch.nn.DataParallel

否

  

torch.nn.parallel.DistributedDataParallel

是

  

torch.nn.parallel.DistributedDataParallel.join

是

  

torch.nn.parallel.DistributedDataParallel.join_hook

是

  

torch.nn.parallel.DistributedDataParallel.no_sync

是

  

torch.nn.parallel.DistributedDataParallel.register_comm_hook

是

  

torch.nn.utils.clip_grad_norm_

     

torch.nn.utils.clip_grad_value_

是

支持bf16,fp32,fp16

torch.nn.utils.parameters_to_vector

     

torch.nn.utils.vector_to_parameters

是

支持bf16,fp16,fp32,fp64,complex64

torch.nn.utils.prune.BasePruningMethod

是

  

torch.nn.utils.prune.BasePruningMethod.apply

是

  

torch.nn.utils.prune.BasePruningMethod.apply_mask

是

支持fp32

torch.nn.utils.prune.BasePruningMethod.abstract compute_mask

是

  

torch.nn.utils.prune.BasePruningMethod.prune

是

支持fp32

torch.nn.utils.prune.BasePruningMethod.remove

是

支持fp32

torch.nn.utils.prune.PruningContainer

是

  

torch.nn.utils.prune.PruningContainer.add_pruning_method

是

  

torch.nn.utils.prune.PruningContainer.apply

是

  

torch.nn.utils.prune.PruningContainer.apply_mask

是

  

torch.nn.utils.prune.PruningContainer.compute_mask

是

支持fp32

torch.nn.utils.prune.PruningContainer.prune

是

支持fp32

torch.nn.utils.prune.PruningContainer.remove

是

支持fp32

torch.nn.utils.prune.Identity

是

支持fp32

torch.nn.utils.prune.Identity.apply

是

支持fp32

torch.nn.utils.prune.Identity.apply_mask

是

支持fp32

torch.nn.utils.prune.Identity.prune

是

支持fp32

torch.nn.utils.prune.Identity.remove

是

支持fp32

torch.nn.utils.prune.RandomUnstructured

是

支持fp32

torch.nn.utils.prune.RandomUnstructured.apply

是

支持fp32

torch.nn.utils.prune.RandomUnstructured.apply_mask

是

支持fp32

torch.nn.utils.prune.RandomUnstructured.prune

是

支持fp32

torch.nn.utils.prune.RandomUnstructured.remove

是

  

torch.nn.utils.prune.L1Unstructured

是

支持fp32

torch.nn.utils.prune.L1Unstructured.apply

是

支持fp32

torch.nn.utils.prune.L1Unstructured.apply_mask

是

支持fp32

torch.nn.utils.prune.L1Unstructured.prune

是

支持fp32

torch.nn.utils.prune.L1Unstructured.remove

是

支持fp32

torch.nn.utils.prune.RandomStructured

是

支持fp32

torch.nn.utils.prune.RandomStructured.apply

是

支持fp32

torch.nn.utils.prune.RandomStructured.apply_mask

是

支持fp32

torch.nn.utils.prune.RandomStructured.compute_mask

是

支持fp32

torch.nn.utils.prune.RandomStructured.prune

是

  

torch.nn.utils.prune.RandomStructured.remove

是

  

torch.nn.utils.prune.LnStructured

是

支持fp32

torch.nn.utils.prune.LnStructured.apply

是

支持fp32

torch.nn.utils.prune.LnStructured.apply_mask

是

支持fp32

torch.nn.utils.prune.LnStructured.compute_mask

是

支持fp32

torch.nn.utils.prune.LnStructured.prune

是

支持fp32

torch.nn.utils.prune.LnStructured.remove

是

支持fp32

torch.nn.utils.prune.CustomFromMask

是

支持int64

torch.nn.utils.prune.CustomFromMask.apply

是

支持int64

torch.nn.utils.prune.CustomFromMask.apply_mask

是

  

torch.nn.utils.prune.CustomFromMask.prune

是

  

torch.nn.utils.prune.CustomFromMask.remove

是

  

torch.nn.utils.prune.Identity

是

支持fp32

torch.nn.utils.prune.Identity.apply

是

支持fp32

torch.nn.utils.prune.Identity.apply_mask

是

支持fp32

torch.nn.utils.prune.Identity.prune

是

支持fp32

torch.nn.utils.prune.Identity.remove

是

支持fp32

torch.nn.utils.prune.random_unstructured

是

  

torch.nn.utils.prune.l1_unstructured

是

  

torch.nn.utils.prune.random_structured

是

  

torch.nn.utils.prune.ln_structured

是

  

torch.nn.utils.prune.global_unstructured

是

  

torch.nn.utils.prune.custom_from_mask

是

支持int64

torch.nn.utils.prune.remove

是

  

torch.nn.utils.prune.is_pruned

是

  

torch.nn.utils.weight_norm

     

torch.nn.utils.remove_weight_norm

     

torch.nn.utils.spectral_norm

是

  

torch.nn.utils.remove_spectral_norm

     

torch.nn.utils.skip_init

否

  

torch.nn.utils.parametrizations.orthogonal

否

  

torch.nn.utils.parametrizations.spectral_norm

否

  

torch.nn.utils.parametrize.register_parametrization

否

  

torch.nn.utils.parametrize.remove_parametrizations

否

  

torch.nn.utils.parametrize.cached

否

  

torch.nn.utils.parametrize.is_parametrized

否

  

torch.nn.utils.parametrize.ParametrizationList

是

  

torch.nn.utils.parametrize.ParametrizationList.right_inverse

是

支持fp32

torch.nn.utils.rnn.PackedSequence

是

支持fp32,int64

torch.nn.utils.rnn.PackedSequence.count

是

支持fp32

torch.nn.utils.rnn.PackedSequence.index

是

支持fp32

torch.nn.utils.rnn.PackedSequence.is_pinned

是

  

torch.nn.utils.rnn.PackedSequence.to

是

支持fp32,int64

torch.nn.utils.rnn.pack_padded_sequence

     

torch.nn.utils.rnn.pad_packed_sequence

     

torch.nn.utils.rnn.pad_sequence

是

支持fp32,fp16

torch.nn.utils.rnn.pack_sequence

     

torch.nn.Flatten

是

支持bf16,fp32,fp16,uint8,int8,int16,int32,int64,bool,complex64,complex128

torch.nn.Unflatten

是

支持fp16,fp32,fp64,int8,int16,uint8,int32,int64,bool

torch.nn.modules.lazy.LazyModuleMixin

是

支持fp32

torch.nn.modules.lazy.LazyModuleMixin.has_uninitialized_params

是

支持fp32

torch.nn.modules.lazy.LazyModuleMixin.initialize_parameters

是

支持fp32