ocr 模型无法导出 FP16 onnx

[zepingWww]

• 系统环境/System Environment：

CentOS Linux 7

NVIDIA-SMI 515.43.04 Driver Version: 515.43.04 CUDA Version: 11.7

• 运行指令/Command Code：

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
from unittest import TestCase

os.environ["FLAGS_allocator_strategy"] = 'auto_growth'

import paddle
# 安装 2.6.0 版本的 ppocr, commit id:4b275f1857dfd44524db0860c6aa9bfcb75a7a81
from ppocr.modeling.architectures import build_model

det_model_config = {'model_type': 'det', 
                    'algorithm': 'DB', 
                    'Transform': None, 
                    'Backbone': {'name': 'ResNet_vd', 'layers': 18, 'disable_se': True}, 
                    'Neck': {'name': 'DBFPN', 'out_channels': 256}, 
                    'Head': {'name': 'DBHead', 'k': 50}}
cls_model_config = {'model_type': 'cls', 
                    'algorithm': 'CLS', 
                    'Transform': None, 
                    'Backbone': {'name': 'MobileNetV3', 'scale': 0.35, 'model_name': 'small'}, 
                    'Neck': None, 
                    'Head': {'name': 'ClsHead', 'class_dim': 2}}
rec_model_config = {'model_type': 'rec', 
                    'algorithm': 'SVTR', 
                    'Transform': None, 
                    'Backbone': {'name': 'MobileNetV1Enhance', 'scale': 0.5, 'last_conv_stride': [1, 2], 'last_pool_type': 'avg'}, 
                    'Head': {'name': 'MultiHead', 'head_list': [{'CTCHead': {'Neck': {'name': 'svtr', 'dims': 64, 'depth': 2, 'hidden_dims': 120, 'use_guide': True}, 'Head': {'fc_decay': 1e-05}}}, {'SARHead': {'enc_dim': 512, 'max_text_length': 25}}], 
                    'out_channels_list': {'CTCLabelDecode': 97, 'SARLabelDecode': 99}}}

input_shape = {
    'det':[2,3,960,960],
    "cls":[2,3,48,192],
    "rec":[2,3,48,320]
}
model_config = {
    'det':det_model_config,
    'cls':cls_model_config,
    'rec':rec_model_config
}

@paddle.no_grad()
def main(model_type:str='det'):
    # build model
    model = build_model(model_config[model_type])
    model.eval()
    dtype = paddle.float16
    model = model.to(dtype=paddle.float16)
    x_spec = paddle.ones(input_shape[model_type], dtype=dtype)
    save_path = f'onnx.save/{model_type}_bs2_fp16'
    paddle.onnx.export(model, save_path, input_spec=[x_spec], opset_version=11)


class TestFP16OCR2Onnx(TestCase):
    def test_det_export_fp16_onnx(self,):
        main(model_dtype='det')

    def test_cls_export_fp16_onnx(self,):
        main(model_dtype='cls')

    def test_rec_export_fp16_onnx(self,):
        main(model_dtype='rec')



if __name__ == '__main__':
  	main('det')
  	main('cls')
    main('rec')

• 完整报错/Complete Error Message：

1. det、cls 模型 cuda 报错 不支持 FP16 计算
   

2. rec 模型 cuda 报错 模型不支持 cast
   大概率是 LSTM 层的问题
   

[zepingWww]

@tink2123 请问一下这个问题能帮忙解答一下吗