torch.nn

torch.nn

表1 torch.nn

序号

API名称

是否支持

1

Parameter

2

UninitializedBuffer

3

UninitializedParameter

Containers

表2 Containers

序号

API名称

是否支持

1

Module

2

ModuleDict

3

ModuleList

4

ParameterDict

5

ParameterList

6

register_module_backward_hook

7

register_module_forward_hook

8

register_module_forward_pre_hook

9

register_module_full_backward_hook

10

Sequential

Convolution Layers

表3 Convolution Layers

序号

API名称

是否支持

1

nn.Conv1d

2

nn.Conv2d

3

nn.Conv3d

4

nn.ConvTranspose1d

5

nn.ConvTranspose2d

6

nn.ConvTranspose3d

7

nn.Fold

8

nn.LazyConv1d

9

nn.LazyConv2d

10

nn.LazyConv3d

11

nn.LazyConvTranspose1d

12

nn.LazyConvTranspose2d

13

nn.LazyConvTranspose3d

14

nn.Unfold

Pooling layers

表4 Pooling layers

序号

API名称

是否支持

1

nn.AdaptiveAvgPool1d

2

nn.AdaptiveAvgPool2d

3

nn.AdaptiveAvgPool3d

是,仅支持D=1,H=1,W=1场景

4

nn.AdaptiveMaxPool1d

5

nn.AdaptiveMaxPool2d

6

nn.AdaptiveMaxPool3d

7

nn.AvgPool1d

8

nn.AvgPool2d

9

nn.AvgPool3d

10

nn.FractionalMaxPool2d

11

nn.FractionalMaxPool3d

12

nn.LPPool1d

13

nn.LPPool2d

14

nn.MaxPool1d

15

nn.MaxPool2d

16

nn.MaxPool3d

17

nn.MaxUnpool1d

18

nn.MaxUnpool2d

19

nn.MaxUnpool3d

Padding Layers

表5 Padding Layers

序号

API名称

是否支持

1

nn.ConstantPad1d

2

nn.ConstantPad2d

3

nn.ConstantPad3d

4

nn.ReflectionPad1d

5

nn.ReflectionPad2d

6

nn.ReflectionPad3d

7

nn.ReplicationPad1d

8

nn.ReplicationPad2d

9

nn.ReplicationPad3d

10

nn.ZeroPad2d

Non-linear Activations (weighted sum, nonlinearity)

表6 Non-linear Activations

序号

API名称

是否支持

1

nn.CELU

2

nn.ELU

3

nn.GELU

4

nn.GLU

5

nn.Hardshrink

6

nn.Hardsigmoid

7

nn.Hardswish

8

nn.Hardtanh

9

nn.LeakyReLU

10

nn.LogSigmoid

11

nn.Mish

12

nn.MultiheadAttention

13

nn.PReLU

14

nn.ReLU

15

nn.ReLU6

16

nn.RReLU

17

nn.SELU

18

nn.Sigmoid

19

nn.SiLU

20

nn.Softplus

21

nn.Softshrink

22

nn.Softsign

23

nn.Tanh

24

nn.Tanhshrink

25

nn.Threshold

Non-linear Activations (other)

表7 Normalization Layers

序号

API名称

是否支持

1

nn.AdaptiveLogSoftmaxWithLoss

2

nn.LogSoftmax

3

nn.Softmax

4

nn.Softmax2d

5

nn.Softmin

Normalization Layers

表8 Non-linear Activations

序号

API名称

是否支持

1

nn.BatchNorm1d

2

nn.BatchNorm2d

3

nn.BatchNorm3d

4

nn.GroupNorm

5

nn.InstanceNorm1d

6

nn.InstanceNorm2d

7

nn.InstanceNorm3d

8

nn.LayerNorm

9

nn.LazyBatchNorm1d

10

nn.LazyBatchNorm2d

11

nn.LazyBatchNorm3d

12

nn.LazyInstanceNorm1d

13

nn.LazyInstanceNorm2d

14

nn.LazyInstanceNorm3d

15

nn.LocalResponseNorm

16

nn.SyncBatchNorm

Recurrent Layers

表9 Recurrent Layers

序号

API名称

是否支持

1

nn.GRU

2

nn.GRUCell

3

nn.LSTM

4

nn.LSTMCell

是,非16对齐场景暂不支持

5

nn.RNN

6

nn.RNNBase

7

nn.RNNCell

Transformer Layers

表10 Transformer Layers

序号

API名称

是否支持

1

nn.Transformer

2

nn.TransformerDecoder

3

nn.TransformerDecoderLayer

4

nn.TransformerEncoder

5

nn.TransformerEncoderLayer

Linear Layers

表11 Linear Layers

序号

API名称

是否支持

1

nn.Bilinear

2

nn.Identity

3

nn.LazyLinear

4

nn.Linear

Dropout Layers

表12 Dropout Layers

序号

API名称

是否支持

1

nn.AlphaDropout

2

nn.Dropout

3

nn.Dropout2d

4

nn.Dropout3d

5

nn.FeatureAlphaDropout

Sparse Layers

表13 Sparse Layers

序号

API名称

是否支持

1

nn.Embedding

2

nn.EmbeddingBag

Distance Functions

表14 Distance Functions

序号

API名称

是否支持

1

nn.CosineSimilarity

2

nn.PairwiseDistance

Loss Functions

表15 Loss Functions

序号

API名称

是否支持

1

nn.BCELoss

2

nn.BCEWithLogitsLoss

3

nn.CosineEmbeddingLoss

4

nn.CrossEntropyLoss

5

nn.CTCLoss

6

nn.GaussianNLLLoss

7

nn.HingeEmbeddingLoss

8

nn.HuberLoss

9

nn.KLDivLoss

10

nn.L1Loss

11

nn.MarginRankingLoss

12

nn.MSELoss

13

nn.MultiLabelMarginLoss

14

nn.MultiLabelSoftMarginLoss

15

nn.MultiMarginLoss

16

nn.NLLLoss

17

nn.PoissonNLLLoss

18

nn.SmoothL1Loss

19

nn.SoftMarginLoss

20

nn.TripletMarginLoss

21

nn.TripletMarginWithDistanceLoss

Vision Layers

表16 Vision Layers

序号

API名称

是否支持

1

nn.PixelShuffle

2

nn.PixelUnshuffle

3

nn.Upsample

4

nn.UpsamplingNearest2d

5

nn.UpsamplingBilinear2d

Shuffle Layers

表17 Shuffle Layers

序号

API名称

是否支持

1

nn.ChannelShuffle

DataParallel Layers (multi-GPU, distributed)

表18 DataParallel Layers

序号

API名称

是否支持

1

nn.DataParallel

2

nn.parallel.DistributedDataParallel

Utilities

表19 Utilities

序号

API名称

是否支持

1

clip_grad_norm_

2

clip_grad_value_

3

nn.Flatten

4

nn.Unflatten

5

nn.utils.rnn.pack_padded_sequence

6

nn.utils.rnn.pack_sequence

7

nn.utils.rnn.PackedSequence

8

nn.utils.rnn.pad_packed_sequence

9

nn.utils.rnn.pad_sequence

10

parameters_to_vector

11

parametrizations.orthogonal

12

parametrizations.spectral_norm

13

parametrize.cached

14

parametrize.is_parametrized

15

parametrize.ParametrizationList

16

parametrize.register_parametrization

17

parametrize.remove_parametrizations

18

prune.BasePruningMethod

19

prune.custom_from_mask

20

prune.CustomFromMask

21

prune.global_unstructured

22

prune.Identity

23

prune.identity

24

prune.is_pruned

25

prune.l1_unstructured

26

prune.L1Unstructured

27

prune.ln_structured

28

prune.LnStructured

29

prune.PruningContainer

30

prune.random_structured

31

prune.random_unstructured

32

prune.RandomStructured

33

prune.RandomUnstructured

34

prune.remove

35

remove_spectral_norm

36

remove_weight_norm

37

skip_init

38

spectral_norm

39

vector_to_parameters

40

weight_norm

Lazy Modules Initialization

表20 Lazy Modules Initialization

序号

API名称

是否支持

1

nn.modules.lazy.LazyModuleMixin