swin_Tranformer使用netron可视化
操作记录和遇到的问题
由于使用netron可视化swin_tiny的时候遇到了一些问题,因此记录一下操作过程和遇到的问题。
准备工作
ubuntu/windows 本人是ubuntu
激活虚拟环境
activate your env :conda activate urenv;
需要用到的package:
- timm :
pip install timm == 0.6.13
- onnx :
conda install onnx
- netron :
pip install netron
- torch:
conda install torch
我的package:
netron=7.6.8=pypi_0
onnx=1.16.0=pypi_0
python=3.10.11=h7a1cb2a_2
pytorch=1.13.1=py3.10_cuda11.6_cudnn8.3.2_0
timm=0.6.13=pypi_0
torch=1.13.1=pypi_0
torchvision=0.14.1=pypi_0
下载模型
从meta官网下载想要的模型
例:swin_tiny的模型
代码
import torch
import timm
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
modelpath=r"/home/usr/.cache/torch/hub/checkpoints/swin_tiny_patch4_window7_224.pth"
# pth模型路径
model = timm.create_model('swin_tiny_patch4_window7_224', pretrained=True)
model.load_state_dict(torch.load(modelpath), strict=False)
# 加载网络结构
model.to(device)
# 模型放到device上
model.eval()
# 将模型置于评估模式
inputs = torch.ones((1, 3, 224,224)).to(device)
# 不考虑预处理,设置一个符合模型输入维度要求的输入。网络输入size 同样放到device上
onnxpath=r"./swin_T.onnx"
# onnx格式的模型保存路径
torch.onnx.export(model, inputs, onnxpath, export_params=True, verbose=True, input_names=['input'], output_names=['output'], opset_version=12)
# 模型格式转换
import onnx
from onnx import shape_inference
model = r'./swin_T.onnx'
#上一步保存好的onnx格式的模型路径
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model)), model)
#增加节点的shape信息
import netron
modelpath=r"./swin_T.onnx"
netron.start(modelpath)
如果想查看timm支持哪些模型
model_list = timm.list_models()
print('\n'.join(model_list))
model_pretrain_list = timm.list_models(pretrained=True)
print('\n'.join(model_pretrain_list))
model_resnet = timm.list_models('*swin*')
print('\n'.join(model_resnet))
结果
遇到的问题
-
1
运行第3块代码时,会在ubuntu上打开firefox然后读取onnx文件,但是不知为何netron可视化不了,意思是需要很长的时间可视化
解决办法:
从系统上下载onnx文件到本地,然后在netron网页端导入该文件进行可视化 -
2
在第一次转成onnx文件后可视化的结果却是ViT的模型
可能原因:
timm库有bug,因为我在使用该虚拟环境时timm已存在,也许有未知的bug没被修复,我用时就出现了
解决办法:
重新安装timm
pip install timm==0.6.13
正确的结果: