This is a great project. Can you provide example code on how to predict images after pruning is completed. I made an error while making the prediction:"AttributeError: Can't get attribute 'C2f_v2' on <module '__main__'"

[qinyaolong123]

This is my predict code:
from ultralytics import YOLO

model = YOLO(
    '/home/PycharmProjects/Torch-Pruning-master/examples/yolov8/runs/detect/step_2_finetune3/weights/yolov8n.pt')
encode_image = cv2.imread(
    "/home/PycharmProjects/FenJiu_KP_gpu/images/pk_cm/val/images/Image_20240513150834969.jpg")

if __name__ == '__main__':
    output = model.predict(encode_image)
    print(output)

This is the error:
Traceback (most recent call last):
File "/home/aliyun/PycharmProjects/Torch-Pruning-master/examples/yolov8/predict.py", line 8, in
model = YOLO(
File "/home/aliyun/anaconda3/envs/yolo/lib/python3.10/site-packages/ultralytics/yolo/engine/model.py", line 107, in init
self._load(model, task)
File "/home/aliyun/anaconda3/envs/yolo/lib/python3.10/site-packages/ultralytics/yolo/engine/model.py", line 156, in _load
self.model, self.ckpt = attempt_load_one_weight(weights)
File "/home/aliyun/anaconda3/envs/yolo/lib/python3.10/site-packages/ultralytics/nn/tasks.py", line 578, in attempt_load_one_weight
ckpt, weight = torch_safe_load(weight) # load ckpt
File "/home/aliyun/anaconda3/envs/yolo/lib/python3.10/site-packages/ultralytics/nn/tasks.py", line 518, in torch_safe_load
return torch.load(file, map_location='cpu'), file # load
File "/home/aliyun/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/serialization.py", line 789, in load
return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
File "/home/aliyun/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/serialization.py", line 1131, in _load
result = unpickler.load()
File "/home/aliyun/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/serialization.py", line 1124, in find_class
return super().find_class(mod_name, name)
AttributeError: Can't get attribute 'C2f_v2' on <module 'main' from '/home/aliyun/PycharmProjects/Torch-Pruning-master/examples/yolov8/predict.py'>

[VainF]

Hi @qinyaolong123, the C2f_v2 module should be imported in your predict.py, since it is defined in another file and pickle serialization will not save those modules for you.