How to convert tranformer ocr model to onnx? about vietocr HOT 4 CLOSED

@abigailng, để export nó mình chia ra 3 cái graph nhỏ: cnn,encoder, decoder.

1. export cnn thì modify function translate() trong vietocr/tool/translate.py như sau:
   def translate(img, model, max_seq_length=128, sos_token=1, eos_token=2): "data: BxCXHxW" model.eval() device = img.device with torch.no_grad(): src = model.cnn(img) torch.onnx.export(model.cnn, img, os.environ['ONNX_PATH']+"/model_cnn.onnx", export_params=True, opset_version=12, do_constant_folding=True, input_names=['img'], output_names=['output'], dynamic_axes={'img': {3: 'lenght'}, 'output': {0: 'channel'}}) .....

2. export encoder và decoder thì modify functions trong vietocr/model/seqmodel/seq2seq.py:
   def forward_encoder(self, src): """ src: timestep x batch_size x channel hidden: batch_size x hid_dim encoder_outputs: src_len x batch_size x hid_dim """ encoder_outputs, hidden = self.encoder(src) torch.onnx.export(self.encoder, src, os.environ['ONNX_PATH']+"/model_encoder.onnx", export_params=True, opset_version=12, do_constant_folding=True, input_names=['src'], output_names=['encoder_outputs', 'hidden'], dynamic_axes={'src':{0: "channel_input"}, 'encoder_outputs': {0: 'channel_output'}}) return (hidden, encoder_outputs)
   `def forward_decoder(self, tgt, memory):
   """
   tgt: timestep x batch_size
   hidden: batch_size x hid_dim
   encouder: src_len x batch_size x hid_dim
   output: batch_size x 1 x vocab_size
   """

    tgt = tgt[-1]
    hidden, encoder_outputs = memory
    output, hidden, _ = self.decoder(tgt, hidden, encoder_outputs)
    torch.onnx.export(self.decoder,
        (tgt, hidden, encoder_outputs),
        os.environ['ONNX_PATH']+"/model_decoder.onnx",
        export_params=True,
        opset_version=12,
        do_constant_folding=True,
        input_names=['tgt', 'hidden', 'encoder_outputs'],
        output_names=['output', 'hidden_out', 'last'],
        dynamic_axes={'encoder_outputs':{0: "channel_input"},
                      'last': {0: 'channel_output'}})
    output = output.unsqueeze(1)
    return output, (hidden, encoder_outputs)`
   

END: việc còn lại là bạn load 3 cái graph này lên. code lại có preprocess và postprocess của tác giả. good luck

from vietocr.

Việc convert sang onnx chạy nhanh hơn 50% so vs chạy trên framework pytorch nha @pbcquoc . Ví dụ trong 1 cái ảnh. khi mình đã chạy cái ảnh đó qua paddleocr detection thì nó sẽ detect ra n cái boxes nhỏ. lúc này nếu dùng padding hay cái gì đó để biến những cái boxes nhỏ này về 1 size r đưa vô 1 batch đổi pass through qua model của Quốc làm giảm performance. Nếu chạy từng cái thì mỗi cái box nhỏ chạy tốn tầm 50-70ms. Trong bào toán của mình 1 ảnh có tới 50-70 cái boxes nhỏ. tính ra model của Quốc sẽ chạy 50-70 lần. end to end tốn tầm 4-6s cho 1 phiếu A4 khi chụp ảnh. Nếu export qua onnx thì thời gian chạy tầm 30-40s cho 1 boxes nhỏ thôi. thời gian đó nhân lơn 40-70 lần là khá nhìu. Kết luận export qua onnx chạy chỉ tầm 2-2.5s cho 1 ảnh. còn nếu chạy pytorch thì tốn 4-6s.

from vietocr.

Hi em,
đầu tiên là em convert sang dạng onnx cũng không nhanh hơn bao nhiêu mà chỉ tiện hơn để deploy thôi ak.
anh chưa em hiểu em xử lý gì nữa mà phải viết lại class nhỉ?

from vietocr.

@thanhpt55 e thấy a đang convert model seq2seq phải ko ạ, a đã convert model transformer của vietocr thử chưa ạ

from vietocr.