剪枝后权重文件导入失败

Question

请教下我用yolov5剪枝后微调后，需要检测一下效果好坏，但是现在出现了权重文件无法导入的情况，不知道问题出在了什么地方？下面是出问题的代码部分，请帮我看看问题在什么地方？

model = DetectPrunedMultiBackend(weights, device=device, dnn=dnn)   # 这是模型导入和重构的类
model = attempt_load_pruned(weights if isinstance(weights, list) else w, map_location='cpu')   在这个类内用attempt_load_pruned()来导入权重
 ckpt = torch.load(attempt_download(w), map_location=map_location)  # 问题出在了这个地方

显示的错误是这样的

 File "D:\anaconda\envs\mmyolo\lib\site-packages\torch\serialization.py", line 875, in find_class
    return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'yolo_pruned'

显示找不到这个模块。我在netron中也能导入权重可视化。

    class UnpicklerWrapper(pickle_module.Unpickler):  # type: ignore[name-defined]
        # from https://stackoverflow.com/questions/13398462/unpickling-python-objects-with-a-changed-module-path/13405732
        # Lets us override the imports that pickle uses when unpickling an object.
        # This is useful for maintaining BC if we change a module path that tensor instantiation relies on.
        def find_class(self, mod_name, name):
            mod_name = load_module_mapping.get(mod_name, mod_name)
            return super().find_class(mod_name, name)

我用1.10版本和1.13版本都测试了，不行。显示同样的错误。

Answer 1

以下答案由GPT-3.5大模型与博主波罗歌共同编写：
根据你提供的信息，问题可能出在权重文件中使用了之前的自定义模块，在加载权重文件时找不到该模块从而导致错误。你可以尝试重新定义该自定义模块并加载权重文件。具体步骤如下：

1. 找出使用了自定义模块的权重文件中的自定义模块名称

在 YOLOv5 中，常见的自定义模块包括 Focus, Conv 和 BottleneckCSP 等。你可以通过查看模型定义文件中的模块定义来确定使用了哪些模块。

1. 重新定义自定义模块

假设使用了 BottleneckCSP 模块，你可以在模型定义文件中重新定义该模块：

class NewBottleneckCSP(nn.Module):
    def __init__(self, c1, c2, n=1, depthwise=False, gw=False, expansion=0.5, sa=False):
        super(NewBottleneckCSP, self).__init__()
        c_ = int(c2 * expansion)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, stride=1)
        self.cv2 = Conv(c1, c_, 1, stride=1)
        self.cv3 = Conv(2 * c_, c2, 1, stride=1)
        self.m = nn.Sequential(*[Bottleneck(c_, c_, depthwise=depthwise, gw=gw, sa=sa) for _ in range(n)])

    def forward(self, x):
        y1 = self.cv2(F.relu(self.cv1(x)))
        y2 = self.m(self.cv3(F.relu(y1)))
        return torch.cat((y1, y2), dim=1)

1. 使用重新定义后的模块加载权重文件

在加载权重文件时，你需要指定 map_location 参数为 "cpu" 并且添加一个 module_mapping 参数，将旧模块名称映射到新模块名称。比如：

from utils.torch_utils import attempt_load_pruned

weights = 'path/to/pruned/weights.pt'
# 将旧模块 "BottleneckCSP" 映射到新模块 "NewBottleneckCSP"
module_mapping = {'BottleneckCSP': 'NewBottleneckCSP'}

model = attempt_load_pruned(weights, map_location='cpu', module_mapping=module_mapping)

以上步骤中 attempt_load_pruned() 函数是 YOLOv5 提供的加载权重文件的便捷函数，可以实现自动转换权重文件、构建模型等功能，你可以根据需要修改其中的参数。
如果我的回答解决了您的问题，请采纳！