基本配置
常见的包
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
1
2
3
4
5
6
2
3
4
5
6
超参数设置
- batch size
- 初始学习率（初始）
- 训练次数（max_epochs）
- GPU 配置
batch_size = 16
# 批次的大小
lr = 1e-4
# 优化器的学习率
max_epochs = 100
1
2
3
4
5
2
3
4
5
除了直接将超参数设置在训练的代码里，我们也可以使用 yaml、json，dict 等文件来存储超参数，这样可以方便后续的调试和修改，这种方式也是常见的深度学习库（mmdetection，Paddledetection，detectron2）和一些 AI Lab 里面比较常见的一种参数设置方式。
基本配置
我们的数据和模型如果没有经过显式指明设备，默认会存储在 CPU 上，为了加速模型的训练，我们需要显式调用 GPU，一般情况下 GPU 的设置有两种常见的方式：
导入包和版本查询
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))
1
2
3
4
5
6
7
2
3
4
5
6
7
可复现性
在硬件设备（CPU、GPU）不同时，完全的可复现性无法保证，即使随机种子相同。但是，在同一个设备上，应该保证可复现性。具体做法是，在程序开始的时候固定 torch 的随机种子，同时也把 numpy 的随机种子固定。
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
3
4
5
6
2
3
4
5
6
显卡设置
如果只需要一张显卡
# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
1
2
2
如果需要指定多张显卡，比如 0，1 号显卡
PyTorch会在第一次导入时缓存CUDA状态 如果设置CUDA_VISIBLE_DEVICES在导入torch之后，设置会失效
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' # 必须在导入torch前设置！
import torch
1
2
3
2
3
也可以在命令行运行代码时设置显卡
CUDA_VISIBLE_DEVICES=0,1 python train.py
1
清除显存
torch.cuda.empty_cache()
1
也可以使用在命令行重置 GPU 的指令
nvidia-smi --gpu-reset -i [gpu_id]
1
上次更新: 2025/06/25, 11:25:50