mmsegmentation导出onnx模型的问题

看到了导出的onnx自动执行了sofmax操作，不想要这个操作

修改如下代码：改为

在encoder_decoder_mapy.py

def simple_test(self, img, img_meta, rescale=True):"""Simple test with single image."""# MOD# seg_logit = self.inference(img, img_meta, rescale)seg_logit, maps = self.inference(img, img_meta, rescale)if torch.onnx.is_in_onnx_export():# our inference backend only support 4D outputseg_pred = seg_logit.unsqueeze(0)return seg_predseg_pred = seg_logit.argmax(dim=1)seg_pred = seg_pred.cpu().numpy()# unravel batch dimseg_pred = list(seg_pred)# MOD# return seg_predreturn seg_pred, maps

torch.onnx.is_in_onnx_export()是判断是否是onnx输出流，将sofmax操作放到后面就可以了

mmsegmentation导出onnx模型的问题相关推荐

torch量化训练后导出onnx模型并运行
1. 依赖项 import io import numpy as npfrom torch import nn import torch.utils.model_zoo as model_zoo im ...
【地平线开发板模型转换】将pytorch生成的onnx模型转换成.bin模型
文章目录 1 获取onnx模型 2 启动docker容器 3 onnx模型检查 3.1 为什么要检查? 3.2 如何操作 4 图像数据预处理 4.1 一些问题的思考 4.2 图片挑选与放置 4.2 使 ...
pytorch将pt模型转onnx模型
pytorch将pt模型转onnx模型一导出ONNX模型 torch.onnx.export( model, # 要导出的模型 args, # 模型的输入参数,输入参数只需满足shape正确 on ...
YOLOv5导出onnx、TrensorRT部署（LINUX）
提示:文章写完后,目录可以自动生成,如何生成可参考右边的帮助文档文章目录前言一.版本声明二.实现步骤 1.训练一个目标检测模型 2.导出onnx模型 3.Netron可视化 4.编译成trtm ...
pyTorch入门（四）——导出Minist模型，C++ OpenCV DNN进行识别
学更好的别人, 做更好的自己. --<微卡智享> 本文长度为2548字,预计阅读8分钟前言前三章介绍了pyTorch训练的相关,我们也保存模型成功了,今天这篇就是使用C++ OpenC ...
利用TensorRT实现神经网络提速(读取ONNX模型并运行)
前些天发现了一个巨牛的人工智能学习网站,通俗易懂,风趣幽默,忍不住分享一下给大家.点击跳转到网站. 前言这篇文章接着上一篇继续讲解如何具体使用TensorRT. 在之前已经写到过一篇去介绍什么是Te ...
onnx模型如何增加或者去除里面node，即修改图方法
有时候我们通过pytorch导出onnx模型,需要修改一下onnx的图结构,怎么修改呢? 下面两个Python实例提供了修改思路. Changing the graph is easier than ...
【yolov5】pytorch模型导出为onnx模型
博主想拿官网的yolov5训练好pt模型,然后转换成rknn模型,然后在瑞芯微开发板上调用模型检测.但是官网的版本对npu不友好,所以采用改进结构的版本: 将Focus层改成Conv层将Swish激 ...
pytorch导出onnx格式模型时，不固定输入输出维度
Pytorch模型转换为onnx格式模型后,模型的输入.输出维度跟转换模型时,用的dummy_input的维度有关系,属于固定尺寸的输入与输出.可以采用以下代码修改onnx模型的输入输出维度: imp ...

mmsegmentation导出onnx模型的问题

mmsegmentation导出onnx模型的问题相关推荐

最新文章

热门文章