torch.nn

torch.nn

表1 torch.nn

序号

API名称

是否支持

1

Parameter

2

UninitializedParameter

3

UninitializedBuffer

Containers

表2 Containers

序号

API名称

是否支持

1

Module

2

Sequential

3

ModuleList

4

ModuleDict

5

ParameterList

6

ParameterDict

7

register_module_forward_pre_hook

8

register_module_forward_hook

9

register_module_backward_hook

10

register_module_full_backward_hook

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.LazyConv1d

8

nn.LazyConv2d

9

nn.LazyConv3d

10

nn.LazyConvTranspose1d

11

nn.LazyConvTranspose2d

12

nn.LazyConvTranspose3d

13

nn.Unfold

14

nn.Fold

Pooling layers

表4 Pooling layers

序号

API名称

是否支持

1

nn.MaxPool1d

2

nn.MaxPool2d

3

nn.MaxPool3d

4

nn.MaxUnpool1d

5

nn.MaxUnpool2d

6

nn.MaxUnpool3d

7

nn.AvgPool1d

8

nn.AvgPool2d

9

nn.AvgPool3d

10

nn.FractionalMaxPool2d

11

nn.FractionalMaxPool3d

12

nn.LPPool1d

13

nn.LPPool2d

14

nn.AdaptiveMaxPool1d

15

nn.AdaptiveMaxPool2d

16

nn.AdaptiveMaxPool3d

17

nn.AdaptiveAvgPool1d

18

nn.AdaptiveAvgPool2d

19

nn.AdaptiveAvgPool3d

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

Padding Layers

表5 Padding Layers

序号

API名称

是否支持

1

nn.ReflectionPad1d

2

nn.ReflectionPad2d

3

nn.ReflectionPad3d

4

nn.ReplicationPad1d

5

nn.ReplicationPad2d

6

nn.ReplicationPad3d

7

nn.ZeroPad2d

8

nn.ConstantPad1d

9

nn.ConstantPad2d

10

nn.ConstantPad3d

Non-linear Activations (weighted sum, nonlinearity)

表6 Non-linear Activations

序号

API名称

是否支持

1

nn.ELU

2

nn.Hardshrink

3

nn.Hardsigmoid

4

nn.Hardtanh

5

nn.Hardswish

6

nn.LeakyReLU

7

nn.LogSigmoid

8

nn.MultiheadAttention

9

nn.PReLU

10

nn.ReLU

11

nn.ReLU6

12

nn.RReLU

13

nn.SELU

14

nn.CELU

15

nn.GELU

16

nn.Sigmoid

17

nn.SiLU

18

nn.Mish

19

nn.Softplus

20

nn.Softshrink

21

nn.Softsign

22

nn.Tanh

23

nn.Tanhshrink

24

nn.Threshold

25

nn.GLU

Non-linear Activations (other)

表7 Normalization Layers

序号

API名称

是否支持

1

nn.Softmin

2

nn.Softmax

3

nn.Softmax2d

4

nn.LogSoftmax

5

nn.AdaptiveLogSoftmaxWithLoss

Normalization Layers

表8 Non-linear Activations

序号

API名称

是否支持

1

nn.BatchNorm1d

2

nn.BatchNorm2d

3

nn.BatchNorm3d

4

nn.LazyBatchNorm1d

5

nn.LazyBatchNorm2d

6

nn.LazyBatchNorm3d

7

nn.GroupNorm

8

nn.SyncBatchNorm

9

nn.InstanceNorm1d

10

nn.InstanceNorm2d

11

nn.InstanceNorm3d

12

nn.LazyInstanceNorm1d

13

nn.LazyInstanceNorm2d

14

nn.LazyInstanceNorm3d

15

nn.LayerNorm

16

nn.LocalResponseNorm

Recurrent Layers

表9 Recurrent Layers

序号

API名称

是否支持

1

nn.RNNBase

2

nn.RNN

3

nn.LSTM

4

nn.GRU

5

nn.RNNCell

6

nn.LSTMCell

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

7

nn.GRUCell

Transformer Layers

表10 Transformer Layers

序号

API名称

是否支持

1

nn.Transformer

2

nn.TransformerEncoder

3

nn.TransformerDecoder

4

nn.TransformerEncoderLayer

5

nn.TransformerDecoderLayer

Linear Layers

表11 Linear Layers

序号

API名称

是否支持

1

nn.Identity

2

nn.Linear

3

nn.Bilinear

4

nn.LazyLinear

Dropout Layers

表12 Dropout Layers

序号

API名称

是否支持

1

nn.Dropout

2

nn.Dropout2d

3

nn.Dropout3d

4

nn.AlphaDropout

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.L1Loss

2

nn.MSELoss

3

nn.CrossEntropyLoss

4

nn.CTCLoss

5

nn.NLLLoss

6

nn.PoissonNLLLoss

7

nn.GaussianNLLLoss

8

nn.KLDivLoss

9

nn.BCELoss

10

nn.BCEWithLogitsLoss

11

nn.MarginRankingLoss

12

nn.HingeEmbeddingLoss

13

nn.MultiLabelMarginLoss

14

nn.HuberLoss

15

nn.SmoothL1Loss

16

nn.SoftMarginLoss

17

nn.MultiLabelSoftMarginLoss

18

nn.CosineEmbeddingLoss

19

nn.MultiMarginLoss

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

parameters_to_vector

4

vector_to_parameters

5

prune.BasePruningMethod

6

prune.PruningContainer

7

prune.Identity

8

prune.RandomUnstructured

9

prune.L1Unstructured

10

prune.RandomStructured

11

prune.LnStructured

12

prune.CustomFromMask

13

prune.identity

14

prune.random_unstructured

15

prune.l1_unstructured

16

prune.random_structured

17

prune.ln_structured

18

prune.global_unstructured

19

prune.custom_from_mask

20

prune.remove

21

prune.is_pruned

22

weight_norm

23

remove_weight_norm

24

spectral_norm

25

remove_spectral_norm

26

skip_init

27

parametrizations.orthogonal

28

parametrizations.spectral_norm

29

parametrize.register_parametrization

30

parametrize.remove_parametrizations

31

parametrize.cached

32

parametrize.is_parametrized

33

parametrize.ParametrizationList

34

nn.utils.rnn.PackedSequence

35

nn.utils.rnn.pack_padded_sequence

36

nn.utils.rnn.pad_packed_sequence

37

nn.utils.rnn.pad_sequence

38

nn.utils.rnn.pack_sequence

39

nn.Flatten

40

nn.Unflatten

Lazy Modules Initialization

表20 Lazy Modules Initialization

序号

API名称

是否支持

1

nn.modules.lazy.LazyModuleMixin