baudm / parseq

Scene Text Recognition with Permuted Autoregressive Sequence Models (ECCV 2022)
https://huggingface.co/spaces/baudm/PARSeq-OCR
Apache License 2.0
570 stars 125 forks source link

How to export model .ckpt to tensorrt cuda #153

Open sh1man999 opened 22 hours ago

sh1man999 commented 22 hours ago

Does anyone have a ready script for exporting a model to tensorrt?

sh1man999 commented 21 hours ago

I don't know if it's right ?

def to_onnx(model_path, device="cuda"):
    import torch
    # Load the checkpoint
    parseq = PARSeq.load_from_checkpoint(model_path)
    parseq.refine_iters = 0
    parseq.decode_ar = False
    image = torch.rand(1, 3, *parseq.hparams.img_size)  # (1, 3, 32, 128) by default
    parseq = parseq.to(device).eval()
    # To ONNX
    parseq.to_onnx('parseq.onnx', image,  do_constant_folding=True, opset_version=14)

numpy 1.26.4 onnx 1.17.0 onnxruntime-gpu 1.19.2 timm 0.9.16 torch 2.3.0+cu121 torchvision 0.18.0+cu121

I'm getting errors:

C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\__init__.py:1559: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert condition, message
C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\timm\models\vision_transformer.py:91: UserWarning: 1Torch was not compiled with flash attention. (Triggered internally at ..\aten\src\ATen\native\transformers\cuda\sdp_utils.cpp:455.)
  x = F.scaled_dot_product_attention(
Traceback (most recent call last):
  File "D:\data\Piton\neural\lpr_engine\src\scripts\convert_to_onnx.py", line 4, in <module>
    ocr_to_onnx("D:/data/Piton/neural/parseq/models/99.7653_99.9614.ckpt")
  File "D:\data\Piton\neural\lpr_engine\src\lpr_engine\ocr\parseq\utils.py", line 14, in to_onnx
    parseq.to_onnx('parseq.onnx', image,  do_constant_folding=True, opset_version=14) 
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\pytorch_lightning\core\module.py", line 1918, in to_onnx
    torch.onnx.export(self, input_sample, file_path, **kwargs)
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\onnx\utils.py", line 516, in export
    _export(
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\onnx\utils.py", line 1612, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\onnx\utils.py", line 1134, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\onnx\utils.py", line 1010, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\onnx\utils.py", line 914, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\jit\_trace.py", line 1310, in _get_trace_graph
    outs = ONNXTracedModule(
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\nn\modules\module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\nn\modules\module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\user\AppData\Local\pypoetry\Cache\virtualenvs\lpr-engine-Ko3JTxdn-py3.10\lib\site-packages\torch\jit\_trace.py", line 138, in forward
    graph, out = torch._C._create_graph_by_tracing(
RuntimeError: 0 INTERNAL ASSERT FAILED at "..\\torch\\csrc\\jit\\ir\\alias_analysis.cpp":621, please report a bug to PyTorch. We don't have an op for aten::full but it isn't a special case.  Argument types: int[], bool, int, NoneType, Device, bool, 

Candidates:
    aten::full.names(int[] size, Scalar fill_value, *, str[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
    aten::full(SymInt[] size, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
    aten::full.names_out(int[] size, Scalar fill_value, *, str[]? names, Tensor(a!) out) -> Tensor(a!)
    aten::full.out(SymInt[] size, Scalar fill_value, *, Tensor(a!) out) -> Tensor(a!)