【Pytorch】服务器Slurm上单/多GPU模型训练及保存加载

一、多GPU模型训练

1、Pytorch多GPU模型训练原理

用简单的话描述一下，以8个GPU为例。首先将模型放到主GPU上，并该模型在其余7个GPU上都复制一份；接着一个batch_size为64的数据传进来时，数据会被分为8份（每份的batch_size为8），这8份数据分别放到8个GPU上；8个GPU分别计算loss，最后loss会被汇总到主GPU上进行反向传播。

2、Pytorch实现

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

CUDA_DEVICES = 0, 1, 2, 3, 4, 5, 6, 7
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2, 3, 4, 5, 6, 7"

if torch.cuda.device_count() > 1:
    net = torch.nn.DataParallel(net, device_ids=CUDA_DEVICES)
    cudnn.benchmark = True
net = net.to(device)

这里的的关键就是设置CUDA_DEVICES和os.environ["CUDA_VISIBLE_DEVICES"]，以及net = torch.nn.DataParallel(net, device_ids=CUDA_DEVICES)。
其他就跟单GPU一样可以愉快地训练了。

二、多GPU模型保存和加载

众所周知，多GPU训练的模型和单GPU训练模型的区别在于多了一层module，因此有的说要在保存的时候

torch.save({
            'epoch': epoch,
            'state_dict': net.module.state_dict(),
            'optimizer': optimizer.state_dict(),
            ...
            }, PATH)

如果保存时没有加上module，加载时就要建一个新的字典来读取参数

state_dict = torch.load('checkpoint.pt')  # 模型可以保存为pth文件，也可以为pt文件。
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:]  # remove `module.`，表面从第7个key值字符取到最后一个字符，正好去掉了module.
    new_state_dict[name] = v #新字典的key值对应的value为一一对应的值。 
# load params
model.load_state_dict(new_state_dict) # 重新加载这个模型。

反正我是通通失败了…报错

RuntimeError: Error(s) in loading state_dict for DataParallel:
    Unexpected running stats buffer(s) "module. ...........................

直到我看到了博客 的第3点！统统不要，统统不要，只要在加载模型之前nn.DataParallel了，保存和加载同单GPU训练时一样的。
保存

torch.save({
            'epoch': epoch,
            'state_dict': net.state_dict(),
            'optimizer': optimizer.state_dict(),
            ...
            }, PATH)

加载

checkpoint = torch.load(CHECKPOINT_PATH)
net.load_state_dict(checkpoint['state_dict'])

注：当然也与可能是因为服务器上用了srun命令，这样是可以的。Anyway，前两种方式都值得一试。

三、Slurm命令使用

服务器上用的是Slurm管理系统，就说一下srun命令

srun --partition=XXX
     --mpi=pmi2
     --gres=gpu:8
     -n1
     --ntasks-per-node=1
     --job-name=TEST
     --kill-on-bad-exit=1
     python main.py 

参数说明

• –gres=gpu:8：使用的GPU数
• –ntasks-per-node=1：单个节点最多的进程个数
• -n1：作业总的进程数。首先会根据-n和–ntasks-per-node来决定分配多少个节点，再根据–gres决定每个节点使用的GPU数

四、代码

详细代码release在github上，关于cifar10的训练，支持单/多GPU训练以及模型加载继续训练。

想要用单/多GPU训练，只要修改可见的GPU

CUDA_DEVICES = 0, 1, 2, 3, 4, 5, 6, 7
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2, 3, 4, 5, 6, 7"

以及srun中的所述三个参数

srun --partition=XXX
     --mpi=pmi2
     --gres=gpu:8
     -n1
     --ntasks-per-node=1
     --job-name=TEST
     --kill-on-bad-exit=1
     python main.py 

想要加载模型继续训练，在代码中修改模型所在文件夹

fold = 'XXXXX' 

然后加入--resume即可

python main.py --resume

Pytorch网站

resources