wdhao / tensorrtcv Goto Github PK

if __name__ == '__main__':

    #!将yolo5s.pt文件保存为yolov5s.pth文件
    model = torch.load(path + 'yolov5s.pt')['model'].float()
    torch.save(model.state_dict(), path + 'yolov5s.pth')  #只保存模型权重参数，不保存模型结构

(py38torch190-smoke) zxl@R9000P:~/mywork/tensorrtCV$ python getWeights_new.py
Traceback (most recent call last):
  File "getWeights_new.py", line 67, in <module>
    model = torch.load(path + 'yolov5s.pt')['model'].float()
  File "/home/zxl/miniconda3/envs/py38torch190-smoke/lib/python3.8/site-packages/torch/serialization.py", line 607, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/zxl/miniconda3/envs/py38torch190-smoke/lib/python3.8/site-packages/torch/serialization.py", line 882, in _load
    result = unpickler.load()
  File "/home/zxl/miniconda3/envs/py38torch190-smoke/lib/python3.8/site-packages/torch/serialization.py", line 875, in find_class
    return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'models'