常见的包

1
2
3
4
5
6
import os
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optimizer

超参数设置

  • batch size
  • 初始学习率（初始）
  • 训练次数（max_epochs）
  • GPU 配置
1
2
3
4
5
batch_size = 16
# 批次的大小
lr = 1e-4
# 优化器的学习率
max_epochs = 100

除了直接将超参数设置在训练的代码里，我们也可以使用 yaml、json，dict 等文件来存储超参数，这样可以方便后续的调试和修改，这种方式也是常见的深度学习库（mmdetection，Paddledetection，detectron2）和一些 AI Lab 里面比较常见的一种参数设置方式。

基本配置

我们的数据和模型如果没有经过显式指明设备，默认会存储在 CPU 上，为了加速模型的训练，我们需要显式调用 GPU，一般情况下 GPU 的设置有两种常见的方式：

导入包和版本查询

1
2
3
4
5
6
7
import torch
import torch.nn as nn
import torchvision
print(torch.__version__)
print(torch.version.cuda)
print(torch.backends.cudnn.version())
print(torch.cuda.get_device_name(0))

可复现性

在硬件设备（CPU、GPU）不同时，完全的可复现性无法保证，即使随机种子相同。但是，在同一个设备上，应该保证可复现性。具体做法是，在程序开始的时候固定 torch 的随机种子，同时也把 numpy 的随机种子固定。

1
2
3
4
5
6
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

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

显卡设置

如果只需要一张显卡

1
2
# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

如果需要指定多张显卡，比如 0，1 号显卡

PyTorch会在第一次导入时缓存CUDA状态
如果设置CUDA_VISIBLE_DEVICES在导入torch之后，设置会失效

1
2
3
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'  # 必须在导入torch前设置！
import torch

也可以在命令行运行代码时设置显卡

1
CUDA_VISIBLE_DEVICES=0,1 python train.py

清除显存

1
torch.cuda.empty_cache()

也可以使用在命令行重置 GPU 的指令

1
nvidia-smi --gpu-reset -i [gpu_id]