Pytorch踩坑记录之GPU训练模型

• 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；

  特别说明 ：

1. 如果，一开始把model按照单GPU的模式载入，然后又打算将model放在多GPUs模式下训练，则应该确保device_ids[0]和device中启用的gpu_id相同；否则报错！！！
2. 当输入模型的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存到指定文件夹，然后下载到本地进行查看。