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 × C × H × W N\times C \times H\times W N×C×H×W,数据范围在[0,1],需要进行转置和规范化。PIL的向量维度为 C × H × N × W C \times H\times N\times W C×H×N×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 × C × H × W N\times C \times H\times W N×C×H×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)
©️2020 CSDN 皮肤主题: 编程工作室 设计师:CSDN官方博客 返回首页