pytorch 小技巧

2 minute read

Published: April 05, 2020

记录一些写 pytorch 时候不会的点。

nn.Module 类定义时不要把网络模型放到某些数据结构中（如 list），否则 .to(device) 找不到网络模型，可以使用 nn.ModuleList()

Variable 模块已与 Tensor 模块合并，直接创建的 Tensor 变量本身带有 requires_grad 属性，不再需要转换为 Variable 变量

在 detach 之前最好先 clone，如a.clone().detach()

创建一个 tensor 指定好 dtype=torch.float 以及 device

不要直接使用 a[:] = x 这样的 inplace 操作，而是使用 torch.where

GPU 使用

指定 GPU

CUDA_VISIBLE_DEVICES="0" python main.py

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'  # 注意修改环境变量后 device index 可能会发生变化

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

use_cuda = True
device = torch.device('cuda:{}'.format(args.gpu) if (torch.cuda.is_available() and use_cuda) else 'cpu')

查看 GPU 是否可用

import torch
torch.cuda.is_available()

使用 cuda

use_cuda = True
device = torch.device('cuda' if (torch.cuda.is_available() and use_cuda) else 'cpu')

保存提取

# 保存整个网络（网络大的时候可能会比较慢）
torch.save(net, 'net.pkl')  
net = torch.load('net.pkl', map_location=device)

# 只保存网络中的参数（速度快, 占内存少）
torch.save(net.state_dict(), 'net_params.pkl')   
net.load_state_dict(torch.load('net_params.pkl'))

tensorboard 使用

安装 tensorboard

pip install tensorboard

记录 tensorboard

from torch.utils.tensorboard import SummaryWriter

# create a summary writer with automatically generated folder name.
writer = SummaryWriter()
# folder location: runs/May04_22-14-54_s-MacBook-Pro.local/

# create a summary writer using the specified folder name.
writer = SummaryWriter(log_dir="runs/experiment")
# folder location: runs/experiment/

# create a summary writer with comment appended.
writer = SummaryWriter(comment="-LR_0.1_BATCH_16")
# folder location: runs/May04_22-14-54_s-MacBook-Pro.local-LR_0.1_BATCH_16/

# scalar
writer.add_scalar(tag, scalar_value, global_step=***)
writer.add_scalars(main_tag, tag_scalar_dict, global_step=***)

# image
# img_tensor: Default is (3, H, W)
writer.add_image(tag, img_tensor, global_step=None, dataformats='CHW')
# img_tensor: Default is (N, 3, H, W)
writer.add_images(tag, img_tensor, global_step=None, dataformats='NCHW')

# graph
writer.add_graph(model, input_to_model=None)

# figure
writer.add_figure(tag, figure, global_step=None)

# close writer
writer.close()

查看 tensorboard

tensorboard --logdir=runs --host=0.0.0.0 --port=8025

torchsummary

安装 torchsummary

pip install torchsummary

使用 torchsummary

from torchsummary import summary
summary(model, input_size, batch_size=-1, device=torch.device('cuda:0'))

tensor 与 numpy 互转

注意：

numpy()、from_numpy() 和 torch.as_tensor() 函数得到的数组共享相同的内存（所以他们之间的转换很快）。
torch.tensor() 函数总是会进行数据拷贝（就会消耗更多的时间和空间），所以返回的数组和原来的数据不再共享内存。

# tensor 转 numpy
a.clone().detach().cpu().numpy()

# numpy 转 tensor
b = torch.from_numpy(a)		# 共享内存

# numpy / list 转 tensor
b = torch.tensor(a)				# 不共享内存
b = torch.as_tensor(a)		# 共享内存

requires_grad

自定义的 Tensor 变量的 requires_grad 属性默认为 False

import torch
a=torch.tensor([1.0])
print(a.requires_grad)  # False
a.requires_grad = True
print(a.requires_grad)  # True

计算内存消耗

基本公式

1 GB = 1024 MB = 1024*1024 KB = 1024*1024*1024 B
1 Byte (字节) = 8 bit (位)

例子

int64: 64 bit = 8 Byte
float32: 32 bit = 4 Byte

import numpy as np

def get_params_size(params_list):
    params_size = sum([np.prod(list(p.size())) for p in params_list]) * 4 / 1024  # float32
    return "{:.0f}KB".format(params_size)

params = list(net.parameters())
params_size = get_params_size(params)
print("params_size: {}".format(params_size))

矩阵乘法

向量点积 torch.dot()
向量矩阵逐元素乘法 torch.mul() 或者 *
二维矩阵相乘 torch.mm()
三维 Batch 矩阵相乘 torch.bmm()
多维数据矩阵相乘（最后两维相乘） torch.matmul()

维度修改

维度压缩 a.squeeze(dim)
维度扩充 a.unsqueeze(dim)
维度重整 a.reshape()
维度重复 a.repeat()
- a 为 (M, N) 则 a.repeat(K, 1) 为 (K * M, N)，而不是 (M * K, N)
- 建议先维度扩充再维度重复