pytorch单卡，多卡模型加载的的问题

1、多卡训练保存的模型 参数会出现 module 关键字
Unexpected key(s) in state_dict: “module.conv1.weight”, “module.bn1.weight”, “module.bn1.bias”, “module.bn1.running_mean”, “module.bn1.running_var”, “module.bn1.num_batches_tracked”, “module.layer1.0.conv1.weight”, “module.layer1.0.bn1.weight”, “module.layer1.0.bn1.bias”, “module.layer1.0.bn1.running_mean”, “module.layer1.0.bn1.running_var”, “module.layer1.0.bn1.num_batches_tracked”, “module.layer1.0.conv2.weight”, “module.layer1.0.bn2.weight”, “module.layer1.0.bn2.bias”, “module.layer1.0.bn2.running_mean”, “module.layer1.0.bn2.running_var”, “module.layer1.0.bn2.num_batches_tracked”, “module.layer1.0.conv3.weight”,
， 如果是多卡保存的模型，用用单卡加载
方法1、
加载模型时出错，因为单卡机器上，没有使用DataParallel来加载模型，所以会出现加载错误。
用这种方法加载模型，会有 module 参数名字，
model = torch.nn.DataParallel(model).cuda()

2、指定一个GPU训练的话

Use CUDA

os.environ[‘CUDA_VISIBLE_DEVICES’] = args.gpu_id
use_cuda = torch.cuda.is_available()
两种方式指定使用cuda

        if args.gpu_id is not None: #输入参数 传递
            torch.cuda.set_device(args.gpu_id)
            model = model.cuda(args.gpu_id)

需要

if args.pretrained:
    print("=> using pre-trained model '{}'".format(args.arch))
    # model = models.__dict__[args.arch](pretrained=True)
    model = models.__dict__[args.arch]()
    model_dict = model.state_dict()
    assert os.path.isfile(args.pretrained), 'Error: no checkpoint directory found!'
    checkpoint = torch.load(args.pretrained)  #  ("sge_resnet50.pth.tar") the author model need to modifify fc
    pretrained_dict = checkpoint['state_dict']

    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in pretrained_dict.items():
        if "fc" in k:  #the author model need to remove fc 1000 param，移除全连接，模型分类输出是10，如果模型不是最后一个全连接，不可以这么移除
            continue
        name = k[7:] # remove `module.` ## 多gpu 训练带moudule默认参数名字,预训练删除
        new_state_dict[name] = v

    model_dict.update(new_state_dict)
    model.load_state_dict(model_dict)

‘’’

判断 GPU，如果是一个，和多个加载不一样

import os
import torch
args.gpu_id=“2,4” ; #指定gpu id
args.cuda = not args.no_cuda and torch.cuda.is_available() #作为是否使用cpu的判定
#配置环境 也可以在运行时临时指定 CUDA_VISIBLE_DEVICES=‘2,6’ Python train.py
os.environ[‘CUDA_VISIBLE_DEVICES’] = args.gpu_id #这里的赋值必须是字符串，list会报错
device_ids=range(torch.cuda.device_count()) #torch.cuda.device_count()=2
#device_ids=[0,1] 这里的0 就是上述指定 2，是主gpu, 1就是7,模型和数据由主gpu分发

if arg.cuda:
model=model.cuda() #这里将模型复制到gpu ,默认是cuda(‘0’)，即转到第一个GPU 2
if len(device_id)>1:
model=torch.nn.DaraParallel(model);#前提是model已经.cuda() 了

单卡，多卡

device_ids = [0,1,2,3]

model = model.cuda(device_ids[0])
model = nn.DataParallel(model, device_ids=device_ids)

optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.001)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

optimizer.zero_grad()     loss.backward()    optimizer.step()
optimizer.module.step()

for param_lr in optimizer.module.param_groups:  # 同样是要加module
        #     param_lr['lr'] = param_lr['lr'] * 0.999
————————————————
版权声明：本文为CSDN博主「CS_lcylmh」的原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_19598705/article/details/80396325