swin_Tranformer使用netron可视化

操作记录和遇到的问题

由于使用netron可视化swin_tiny的时候遇到了一些问题，因此记录一下操作过程和遇到的问题。

准备工作

ubuntu/windows 本人是ubuntu

激活虚拟环境

activate your env :conda activate urenv；

需要用到的package：

1. timm :pip install timm == 0.6.13
2. onnx :conda install onnx
3. netron :pip install netron
4. 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
  正确的结果：