pytorch 加载部分参数,以及多GPU 加载单GPUmodel,pop去除全连接参数,然后加载

            ######### multi gpu  load one gpu need mudule.
            print ("model  load  param###########################")
            pretrained_dict = torch.load("model_ir_se50.pth")   ###
            self.model_dict = self.model.state_dict()  #get the name:value
            param={}
            for k, v in pretrained_dict.items():
                if k[7:] !="module" :
                    param["module."+k] = pretrained_dict[k]

            pretrained_dict = { k: v for k, v in param.items() if k in self.model_dict}
            self.model_dict.update(pretrained_dict)
            self.model.load_state_dict(self.model_dict)


实际上就是字典的操作,那么字典的操作,哪些层不要,打印出层的名字就可以了,例如
模型参数的某些层的权重不要,那么重构一个字典参数就可以了,for k in torch.load("**.pth').keys():   打印出来按照名字删除 键值对,
所以字典的pop删除操作也是可以的, 多卡训练参数多了一个module

pretrain = torch.load("/home/imagenet.pth")
new_state_dict = {}#OrderedDict()
for k,v in pretrain.items():
    if "classifier" in k:  #最后分类层的参数是classeifer ,不需要这个模型参数
        continue
    new_state_dict[k[7:]] = v  #remove `module.`  #模型k 有module 不要
model.load_state_dict(new_state_dict, strict=False)  #strict =False ,模型参数和模型不一致可以加载

 model = DPN(num_init_features=64, k_R=96, G=32, k_sec=(3,4,20,3), inc_sec=(16,32,24,128), num_classes=1,decoder=args.decoder)
 http = {'url': 'http://data.lip6.fr/cadene/pretrainedmodels/dpn92_extra-b040e4a9b.pth'}
 pretrained_dict=model_zoo.load_url(http['url'])
 model_dict = model.state_dict()
 pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}#filter out unnecessary keys 
 model_dict.update(pretrained_dict)
 model.load_state_dict(model_dict)
 model = torch.nn.DataParallel(model).cuda()
简单写法

直接pop 去除全连接参数,然后加载,字典的去除键值对的方法

            ###############加载自己训练的模型
            print ("加载自己训练的模型")
            pretrained_dict = torch.load("model_best.pth.tar")["state_dict"] #保存的模型的全部结果包括全连接,优化器
            for k,v in pretrained_dict.items():
                print (k)
                print (v.shape)
            pretrained_dict.pop('_fc.weight') # 加载的参数直接删除全连接层的参数,
            pretrained_dict.pop('_fc.bias')
            model.load_state_dict(pretrained_dict, strict=False)

参数 打印结果,有分类输出层

_blocks.15._se_reduce.bias
torch.Size([48])
_blocks.15._se_expand.weight
torch.Size([1152, 48, 1, 1])
_blocks.15._se_expand.bias
torch.Size([1152])
_blocks.15._project_conv.weight
torch.Size([320, 1152, 1, 1])
_blocks.15._bn2.weight
torch.Size([320])
_blocks.15._bn2.bias
torch.Size([320])
_blocks.15._bn2.running_mean
torch.Size([320])
_blocks.15._bn2.running_var
torch.Size([320])
_blocks.15._bn2.num_batches_tracked
torch.Size([])
_conv_head.weight
torch.Size([1280, 320, 1, 1])
_bn1.weight
torch.Size([1280])
_bn1.bias
torch.Size([1280])
_bn1.running_mean
torch.Size([1280])
_bn1.running_var
torch.Size([1280])
_bn1.num_batches_tracked
torch.Size([])
_fc.weight
torch.Size([115, 1280])
_fc.bias
torch.Size([115])
Using image size 224

  • 2
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
### 回答1: 在PyTorch加载预训练模型时,如果不想使用多GPU,可以通过以下步骤进行操作。 首先,我们需要导入必要的库和模块: ```python import torch import torchvision.models as models ``` 接下来,我们可以选择要加载的预训练模型。这里以ResNet-50为例: ```python model = models.resnet50(pretrained=True) ``` 接下来,我们需要检查当前系统中的GPU数量。可以使用`torch.cuda.device_count()`来获取GPU的数量。 ```python gpu_count = torch.cuda.device_count() ``` 如果`gpu_count`的值大于1,表示有多个GPU可用。为了确保模型只在GPU上运行,我们需要使用`torch.nn.DataParallel`模块。但是,如果不想使用多GPU,则需要取消这个模块的使用。 如果只有一个GPU可用,我们可以简地将模型转移到该GPU上: ```python if gpu_count == 1: model = model.cuda() ``` 如果有多个GPU可用,我们可以使用`torch.nn.DataParallel`模块来加载预训练模型,并将模型转移到主GPU上: ```python if gpu_count > 1: model = torch.nn.DataParallel(model).cuda() ``` 最后,我们可以使用加载的预训练模型进行后续操作。 总之,为了不使用多GPU进行模型加载,我们需要通过检查GPU数量,并相应地转移到GPU或使用`torch.nn.DataParallel`模块转移到主GPU上。这样可以确保模型在GPU上运行而不会使用多GPU。 ### 回答2: 在PyTorch中,加载预训练模型时,如果不希望使用多个GPU上的module,可以通过以下步骤实现: 首先,使用torch.load()函数加载预训练模型的权重和参数,例如: ```python model_state_dict = torch.load('pretrained_model.pth') ``` 然后,在加载模型之前,可以通过修改model_state_dict来删除原模型中包含的多GPU相关的module前缀。默认情况下,PyTorch在保存模型时,会自动添加"module."前缀来标识使用多GPU。例如,如果模型原本的键名为"module.conv1.weight",则可以通过以下代码提取模型参数的键名: ```python new_model_state_dict = {} for k, v in model_state_dict.items(): name = k[7:] # 去掉"module."前缀 new_model_state_dict[name] = v ``` 接下来,创建新的模型实例,并将修改后的参数加载到该实例中: ```python model = YourModel() # 创建新的模型实例 model.load_state_dict(new_model_state_dict) # 加载修改后的模型参数 ``` 注意,这里的YourModel()应该是与预训练模型相同的模型类实例化得到的对象。 通过这样的处理,就可以将预训练模型加载GPU上的模型中,而无需考虑多GPU的module问题。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值