Pytorch常用小函数

参考链接：机器之心

环境方面

查看Pytorch以及GPU相关信息：

torch.__version__               # PyTorch version
torch.version.cuda              # Corresponding CUDA version
torch.backends.cudnn.version()  # Corresponding cuDNN version
torch.cuda.get_device_name(0)   # GPU type，查看GPU型号，可以为0,1,2根据编号

固定随机种子：

np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

设定可见的GPU：

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

判断CUDA支持：

torch.cuda.is_available()

设置为 cuDNN benchmark 模式：
Benchmark 模式会提升计算速度，但是由于计算中有随机性，每次网络前馈结果略有差异。

torch.backends.cudnn.benchmark = True

如果想要避免这种结果波动，需要关闭benchmark 模式，并设置

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

张量操作

张量的基本信息：

tensor.type()   # Data type
tensor.size()   # Shape of the tensor. It is a subclass of Python tuple
tensor.dim()    # Number of dimensions.

数据类型转换：

# Set default tensor type. Float in PyTorch is much faster than double.
torch.set_default_tensor_type(torch.FloatTensor)

# Type convertions.
tensor = tensor.cuda()
tensor = tensor.cpu()
tensor = tensor.float()
tensor = tensor.long()

torch.Tensor 与 np.ndarray 转换

# torch.Tensor -> np.ndarray.
ndarray = tensor.cpu().numpy() # 共享内存

# np.ndarray -> torch.Tensor.
tensor = torch.from_numpy(ndarray).float() # 共享内存
tensor = torch.from_numpy(ndarray.copy()).float()  # If ndarray has negative stride， 不共享内存

negative stride表示np.ndarray执行了torch.Tensor没有的方法就会报错，使用.copy()创建副本之后二者相互独立。
torch.Tensor 与 PIL.Image 转换：注意Tensor的张量维度是$N\times C\times H\times W$,数据范围在[0,1]，需要进行转置和规范化。PIL的向量维度为$C\times H\times N\times W$,数据范围在[0,255]。

# torch.Tensor -> PIL.Image.
image = PIL.Image.fromarray(torch.clamp(tensor * 255, min=0, max=255
    ).byte().permute(1, 2, 0).cpu().numpy())
image = torchvision.transforms.functional.to_pil_image(tensor)  # Equivalently way

# PIL.Image -> torch.Tensor.
tensor = torch.from_numpy(np.asarray(PIL.Image.open(path))
    ).permute(2, 0, 1).float() / 255
tensor = torchvision.transforms.functional.to_tensor(PIL.Image.open(path))  # Equivalently way

np.ndarray 与 PIL.Image 转换：

# np.ndarray -> PIL.Image.
image = PIL.Image.fromarray(ndarray.astypde(np.uint8))

# PIL.Image -> np.ndarray.
ndarray = np.asarray(PIL.Image.open(path))

打乱顺序：

tensor = tensor[torch.randperm(tensor.size(0))]  # Shuffle the first dimension

水平翻转：
PyTorch 不支持 tensor[::-1] 这样的负步长操作，水平翻转可以用张量索引实现。假设张量维度为$N\times C\times H\times W$

tensor = tensor[:, :, :, torch.arange(tensor.size(3) - 1, -1, -1).long()]

复制张量：
提供3种方法对应不同的需求。不在计算图中表明不会计算梯度了。

# Operation                 |  New/Shared memory | Still in computation graph |
tensor.clone()            # |        New         |          Yes               |
tensor.detach()           # |      Shared        |          No                |
tensor.detach.clone()()   # |        New         |          No                |

拼接张量：
注意 torch.cat 和 torch.stack 的区别在于 torch.cat 沿着给定的维度拼接，而 torch.stack 会新增一维。

>>> list = [torch.tensor([[1,2],[3,4]]), torch.tensor([[5,6],[7,8]])]
>>> torch.cat(list, dim=1)
tensor([[1, 2, 5, 6],
        [3, 4, 7, 8]])
>>> torch.stack(list, dim=1)
tensor([[[1, 2],
         [5, 6]],

        [[3, 4],
         [7, 8]]])

矩阵相乘：

# Matrix multiplication: (m*n) * (n*p) -> (m*p).
result = torch.mm(tensor1, tensor2)

# Batch matrix multiplication: (b*m*n) * (b*n*p) -> (b*m*p).
result = torch.bmm(tensor1, tensor2)

# Element-wise multiplication.
result = tensor1 * tensor2
# 或者是
result = torch.mul(tensor1, tensor2)