- Case-1. 确定机器的cuda是否可用、GPU的个数
a) torch.cuda.is_available() # check cuda是否可用
b) nvidia-smi # 窗口 print 显卡信息
c) torch.cuda.device_count() # 统计显卡的个数
d) torch.cuda.current_device() # 返回当前显卡的序号[默认当前显卡index为0]
- Case-2. 采用单GPU训练模型
model = BaseNet() # basenet
device = t.device("cuda:0" if t.cuda.is_available() else "cpu")
model = model.to(device)
备注:这里的device是自定义的,这样做的好处是:它可以灵活指定默认的gpu-index, eg : "cuda:1"
- Case-3. 采用多GPU训练模型
关键命令:torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
参数解释 :device_ids如果不具体指定,则默认启用全部的gpus;
特别说明 :
- 如果,一开始把model按照单GPU的模式载入,然后又打算将model放在多GPUs模式下训练,则应该确保device_ids[0]和device中启用的gpu_id相同;否则报错!!!
- 当输入模型的input-size相同时,使用torch.backends.cudnn.benchmark = True将使得cudnn去寻找最优的资源配置来运行算法,从而加速训练过程;否则,如果input-size存在变化,使用torch.backends.cudnn.benchmark = True则反而拖慢训练速度!!!
官方教程 : tutorial,源码解析 : nn.DataParallel .
'''
step - 1. set the gpu_ids [make sure batch_size > len(gpu_ids)]
step - 2. execute forward/backward on mulitple GPUs
'''
gpu_ids = [0, 1, 2]
parallel_net = torch.nn.DataParallel(model, device_ids=gpu_ids)
- Case-4. 训练过程中调用网络的子模块(以fc层为例)
# single-GPU
model = Net()
output = model.fc(input)
# multiple-GPUs
model = Net()
model = nn.DataParallel(model)
output = model.module.fc(input)
- Case-5. 在远程服务器上调用matplotlib库
import matplotlib.pyplot as plt
plt.switch_backend('agg')
备注:由于远程服务器可能没有图形界面,因此在需要切换plt对应的终端;此外,不建议在代码中直接使用plt.show(),可以直接将figure存到指定文件夹,然后下载到本地进行查看。