遇到的PyTorch API

1. torch.chunk

torch.chunk(input, chunks, dim=0) → List of Tensors

将input tensor划分成特定的块数，每个块都是input tensor的一个视图，最后一个块可能会小一点，因为不能被dim整除。

• input 输入tensor
• chunks 返回多少个块
• dim 沿着哪个维度进行切分

>>> import torch
>>> a = torch.zeros([6, 8 , 2, 2])
>>> x,y = a.chunk(chunks=2,dim=2)
>>> x.shape
torch.Size([6, 8, 1, 2])
>>> y.shape
torch.Size([6, 8, 1, 2])
>>>

2. nn.GroupNorm

BN是Batch维度上的归一化，GN就是Group维度上的归一化。

$$y=\frac{x-\mathrm{E}[x]}{\sqrt{\mathrm{Var}[x]+ϵ}}\ast \gamma +\beta$$

torch.nn.GroupNorm(num_groups: int, num_channels: int, eps: float = 1e-05, affine: bool = True)

• num_groups ：组个数
• num_channels: 输入tensor的channel个数
• eps: 参数默认为1e-05
• affine: 当设置为true，这个模块有可学习的每个通道的仿射参数初始化为1(权重)和0(偏差)。

>>> import torch
>>> a = torch.zeros([3,128,2,2])
>>> import torch.nn as nn
>>> m = nn.GroupNorm(4, 128, 1e-5, True)
>>> m(a).shape
torch.Size([3, 128, 2, 2])
>>>

3. torch.permute

返回维度排列后的原始张量的视图。

>>> x = torch.randn(2, 3, 5)
>>> x.size()
torch.Size([2, 3, 5])
>>> x.permute(2, 0, 1).size()
torch.Size([5, 2, 3])