copying a param of torch.Size([128]) from checkpoint, where the shape is torch.Size([1, 128])解决方式

2024年6月11日 458次阅读

关于torch的在加载bn_inception、x_inception时出现copying a param of torch.Size([128]) from checkpoint, where the shape is torch.Size([1, 128])类型的错误，观察发现可以用过修改加载模型中权重维度解决问题。

例如：bn_inception中都是两维元组的数据转为一维：如（1,128）转为（128）可使用一下程序解决，类似问题可通过微调以下代码解决：

import torch
state_dict = torch.load('~.torch/models/bn_inception-9f5701afb96c8044.pth')
for name, weights in state_dict.items():
    # print(name, weights.size())  可以查看模型中的模型名字和权重维度
    if len(weights.size()) == 2: 判断需要修改维度的条件
        state_dict[name] = weights.squeeze(0)  去掉维度0，把(1,128)转为(128)
       # print(name,weights.squeeze(0).size()) 查看转化后的模型名字和权重维度
torch.save(state_dict, '~.torch/models/bn_inception-9f5701afb96c8044.pth')

 

点赞