pytorch使用多GPU进行训练

首先需要在代码开头注明所使用的GPU序号，比如：

import torch.nn as nn
import os

os.environ["CUDA_VISIBLE_DEVICES"] = '0，1'

对linux系统来说，可以使用

watch -n 0.1 nvidia-smi

来查看服务器上GPU的状态与可用GPU序号。
pytorch多GPU训练有两种方法，DataParallel与DistributedDataParallel，一般来说后者比前者更优一点。pytorch的官网建议使用DistributedDataParallel来代替DataParallel, 据说是因为DistributedDataParallel比DataParallel运行的更快, 然后显存分屏的更加均衡. 而且DistributedDataParallel功能更加强悍, 例如分布式的模型(一个模型太大, 以至于无法放到一个GPU上运行, 需要分开到多个GPU上面执行). 只有DistributedDataParallel支持分布式的模型像单机模型那样可以进行多机多卡的运算.

1.torch.nn.DataParallel

在模型声明的地方修改代码，例如：

    model = Darknet(opt.model_def)
    model.apply(weights_init_normal)
    model = nn.DataParallel(model)
    model = model.to(device)

2.torch.nn.parallel.DistributedDataParallel

因为DistributedDataParallel是支持多机多卡的, 所以这个需要先初始化一下,

torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=0, world_size=1)

之后就和DataParallel类似了：

model = nn.parallel.DistributedDataParallel(model)