Pytorch CookBook

结合网上收集和个人总结，用于以后背记。

检查 PyTorch 版本：

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

固定随机种子

torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

GPU和cudnn

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' #指定GPU
torch.cuda.empty_cache() #清空显存
torch.backends.cudnn.benchmark = True #能够加快计算速度，但每次前馈结果有波动
torch.backends.cudnn.deterministic = True #加上此句，使每次前馈结果相同

torch.Tensor 与 np.ndarray 与PIL.Image的互换

# torch.Tensor <-> np.ndarray
ndarray = tensor.cpu().numpy()
tensor = torch.from_numpy(ndarray).float()
# 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.
image = PIL.Image.fromarray(ndarray.astypde(np.uint8))
ndarray = np.asarray(PIL.Image.open(path))

打乱顺序

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

水平翻转
PyTorch 不支持 tensor[::-1] 这样的负步长操作，水平翻转可以用张量索引实现。

tensor = tensor[:, :, :, torch.arange(tensor.size(3) - 1, -1, -1).long()] #N*D*H*W

复制张量
有三种复制的方式，对应不同的需求。

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

拼接张量

tensor = torch.cat(list_of_tensors, dim=0)  #3个10×5 -> 30x5
tensor = torch.stack(list_of_tensors, dim=0) #3个10×5 -> 3x10x5

独热编码

N = tensor.size(0)
one_hot = torch.zeros(N, num_classes).long()
one_hot.scatter_(dim=1, index=torch.unsqueeze(tensor, dim=1), src=one_hot)

计算模型总参数量

num_parameters = sum(torch.numel(parameter) for parameter in model.parameters())

以较大学习率微调全连接层，以较小学习率微调卷积层

model = torchvision.models.resnet18(pretrained=True)
finetuned_parameters = list(map(id, model.fc.parameters()))
conv_parameters = (p for p in model.parameters() if id(p) not in finetuned_parameters)
parameters = [{'params': conv_parameters, 'lr': 1e-3}, 
              {'params': model.fc.parameters()}]
optimizer = torch.optim.SGD(parameters, lr=1e-2, momentum=0.9, weight_decay=1e-4)

标签平滑

for images, labels in train_loader:
    images, labels = images.cuda(), labels.cuda()
    N = labels.size(0)
    # C is the number of classes.
    smoothed_labels = torch.full(size=(N, C), fill_value=0.1 / (C - 1)).cuda()
    smoothed_labels.scatter_(dim=1, index=torch.unsqueeze(labels, dim=1), value=0.9)

    score = model(images)
    log_prob = torch.nn.functional.log_softmax(score, dim=1)
    loss = -torch.sum(log_prob * smoothed_labels) / N
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

mixup

beta_distribution = torch.distributions.beta.Beta(alpha, alpha)
for images, labels in train_loader:
    images, labels = images.cuda(), labels.cuda()

    # Mixup images.
    lambda_ = beta_distribution.sample([]).item()
    index = torch.randperm(images.size(0)).cuda()
    mixed_images = lambda_ * images + (1 - lambda_) * images[index, :]

    # Mixup loss.    
    scores = model(mixed_images)
    loss = (lambda_ * loss_function(scores, labels) 
            + (1 - lambda_) * loss_function(scores, labels[index]))

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

得到当前学习率

# If there is one global learning rate (which is the common case).
lr = next(iter(optimizer.param_groups))['lr']

# If there are multiple learning rates for different layers.
all_lr = []
for param_group in optimizer.param_groups:
    all_lr.append(param_group['lr'])

注意事项

模型定义

• 建议有参数的层和汇合（pooling）层使用 torch.nn 模块定义，激活函数直接使用torch.nn.functional。torch.nn 模块和 torch.nn.functional 的区别在于，torch.nn模块在计算时底层调用了 torch.nn.functional，但 torch.nn模块包括该层参数，还可以应对训练和测试两种网络状态。
• 使用 torch.nn.functional 时要注意网络状态，如 model(x)前用 model.train() 和 model.eval() 切换网络状态。
• 不需要计算梯度的代码块用 with torch.no_grad() 包含起来。model.eval() 和 torch.no_grad()的区别在于，model.eval() 是将网络切换为测试状态，例如 BN和随机失活（dropout）在训练和测试阶段使用不同的计算方法。torch.no_grad() 是关闭 PyTorch张量的自动求导机制，以减少存储使用和加速计算，得到的结果无法进行 loss.backward()。
• torch.nn.CrossEntropyLoss 的输入不需要经过 Softmax。torch.nn.CrossEntropyLoss等价于 torch.nn.functional.log_softmax + torch.nn.NLLLoss。
  用法： loss = torch.nn.CrossEntropyLoss(); loss(input,target)。
  其中 input:(N,C) 类型torch.float(); target:(N) 类型torch.long() 0≤target[i]≤C−1
• loss.backward() 前用 optimizer.zero_grad() 清除累积梯度。如果需要累积梯度以进行batch_size放大或者多个损失函数相加时可以不清除梯度。optimizer.zero_grad()和 model.zero_grad() 效果一样。

PyTorch 性能与调试

• torch.utils.data.DataLoader 中尽量设置 pin_memory=True，对特别小的数据集如 MNIST 设置pin_memory=False 反而更快一些。num_workers 的设置需要在实验中找到最快的取值。
• 用 del 及时删除不用的中间变量，节约 GPU 存储。
• 使用 inplace 操作可节约 GPU 存储，如

x = torch.nn.functional.relu(x, inplace=True)

• 时常使用 assert tensor.size() == (N, D, H, W) 作为调试手段，确保张量维度和你设想中一致。
• 除了标记 y 外，尽量少使用一维张量，使用 n*1 的二维张量代替，可以避免一些意想不到的一维张量计算结果。
• 统计代码各部分耗时：

with torch.autograd.profiler.profile(enabled=True, use_cuda=False) as profile:
    ...
print(profile)