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 output
        seg_pred = seg_logit.unsqueeze(0)
        return seg_pred
    seg_pred = seg_logit.argmax(dim=1)
    seg_pred = seg_pred.cpu().numpy()
    # unravel batch dim
    seg_pred = list(seg_pred)
    # MOD
    # return seg_pred
    return seg_pred, maps

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

Python中,你可以使用深度学习框架如TensorFlow、PyTorch等将训练好的模型换为ONNX(开放神经网络交换格式)模型。下面是使用这两个框架导出ONNX模型的基本步骤: **TensorFlow** 1. 首先,确保已经安装了`tensorflow`和`onnx`库。 2. 使用模型的 `.save()` 或者 `.export_savedmodel()` 方法保存模型(对于`.h5`文件的Keras模型,可以使用`tf.keras.models.save_model()`)。 3. 导入`tf2onnx`库并加载模型: ```python import tensorflow as tf from tf2onnx import convert_from_tf # 假设你有一个名为model的TensorFlow模型 model = ... # load your TensorFlow model here ``` 4. 模型: ```python onnx_model = convert_from_tf(model, input_names=['input'], output_names=['output']) ``` 5. 将换后的ONNX模型保存到文件: ```python with open('model.onnx', 'wb') as f: f.write(onnx_model.SerializeToString()) ``` **PyTorch** 1. 安装`torch.onnx`库,以及`onnx`库。 2. 将PyTorch模型换为eval模式: ```python model = ... # load your PyTorch model here model.eval() ``` 3. 准备输入数据,并将其换为张量: ```python input_data = ... # prepare your input data input_tensor = torch.from_numpy(input_data).unsqueeze(0) # add batch dimension if needed ``` 4. 使用`.to_onnx()`函数导出ONNX模型: ```python export_path = 'model.onnx' torch.onnx.export(model, input_tensor, export_path, opset_version=11) ``` **相关问题--:** 1. 如何验证导出ONNX模型是否正确? 2. ONNX支持哪些常见的深度学习架构? 3. 如果遇到模型无法换的问题,有哪些解决方法?
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值