xxx.pt is a zip archive(did you mean to use torch.jit.load()?)“

逸梦er

于 2022-09-30 15:50:53 发布

阅读量341

点赞数

CC 4.0 BY-SA版权

文章标签： python

本文链接：https://blog.youkuaiyun.com/yunfeijiang/article/details/127124834

该博客介绍如何使用PyTorch在1.6及以上版本下，批量加载并转换U2NETPP模型为非zip格式，便于进一步操作。通过循环遍历目录下的模型文件，实现状态字典的加载和保存，保持模型结构完整。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

```
import torch import os from model import U2NETPP net = U2NETPP(3,1) #网络实例化 model_dir = 'E:\\model_pp\\' model_dir2 = 'E:\\model_pp\\after\\' model_path = [model_dir + f for f in os.listdir(model_dir) if f.endswith('pth')] for file in model_path: if not os.path.exists(model_dir2): os.makedirs(model_dir2, exist_ok=True) print(file) img_name = file.split(os.sep)[-1] aaa = img_name.split(".") bbb = aaa[0:-1] imidx = bbb[0] state_dict = torch.load(file, map_location='cpu') net.load_state_dict(state_dict) torch.save(net.state_dict(),model_dir2+'u2netpp'+imidx+'.pth', _use_new_zipfile_serialization=False)
```

使用torch1.6及以上版本运行，重新加载保存模型即可

循环将路径下的模型文件全部转换为非zip