ããã«ã¡ã¯ãã€ãããŒã·ã§ã³ã»ã³ã¿ãŒã®å è€ã§ããæ®æ®µã¯ã³ã³ãã¥ãŒã¿ããžã§ã³ã®æè¡éçºãAIã·ã¹ãã ã®æ€èšŒã«åãçµãã§ããŸãã ä»åã¯ææ°çã®PyTorchã䜿ã£ãŠè»œéãªTransformerããŒã¹OCRã¢ãã«ã§ããPARSeq(Permuted Autoregressive Sequence)ãTensorRTã¢ãã«ã«å€æããŠé«éåããåãçµã¿ã«ã€ããŠç޹ä»ããŸãã PARSeqãšã¯ PARSeqã®TensorRTå PyTorch Lightningã«ããã¢ãã«å€æ AutoregressiveãšIterative refinementãTensorRTåã§ããªãåé¡ Autoregressive modeã®TensorRTå TorchDynamoã®æ©å«ããšã Iterative refinementã®TensorRTå è©äŸ¡ ãŸãšã PARSeqãšã¯ PARSeq 1 ã¯Vision Transformer(ViT)ãç¹åŸŽæœåºåšãšããŠçšããæåèªèã¢ãã«ã§ããã以äžã®ç»åã®ãããªæç« çæã®åœ¢ããšã£ãŠããŸãã ãã®ãããªæç« çæã¢ãã«ã§ã¯ããŸãç»åãããŒã¯ã³ã«åå²ãããã®ãTransformer Encoderã§ç¹åŸŽæœåºãããããããšã«Transformer Decoderã§æ¬¡ã®æåããŒã¯ã³ã®äºæž¬ãç¹°ãè¿ããŸããPARSeqã®å Žåã¯æåããŒã¯ã³ã®äºæž¬æ¹æ³ã«ãªãã·ã§ã³ãããã以åã®äºæž¬ãåç
§ããªããïŒæåãã€äºæž¬ãããã®(Autoregressive)ãäžåºŠã«å
šéšã®æåãäºæž¬ãããã®(Non-autoregressive)ãäžåºŠäºæž¬ããæåãå
¥åãçŽããŠæŽç·Žãããã®(Iterative refinement)ã®äžéãã®ãã³ãŒãæŠç¥ããããŸãã PARSeqã®ç¹åŸŽã¯TransformerããŒã¹ã§ãããªããéåžžã«è»œéã§ããç¹ã§ãã Encoderéšåã¯äžè¬çãªViTãšåæ§ã«12å±€ã®Transformerã¬ã€ã€ãŒã§æ§æãããŠããŸãããDecoderéšåã¯ãã£ã1å±€ãããªãã äžè¬çãªVision Language Modelãæ°ååã®ãã©ã¡ãŒã¿ãæ±ããŠããäžæ¹ã§PARSeqã¯æ°åäžãã©ã¡ãŒã¿ã«çãŸã£ãŠããŸãã PARSeqã®TensorRTå ãã®PARSeqã¢ãã«ãããã«é«éåããããã«ãä»åã¯TensorRTã¢ãã«ã«å€æããŸãã TensorRT 2 ã¯ãNVIDIAãæäŸããŠãããã£ãŒãã©ãŒãã³ã°ã¢ãã«ã®æšè«ãé«éåããããã®ããŒã«ã§ãããŸããŸãªAIãã¬ãŒã ã¯ãŒã¯ã察å¿ããŠããå
±éãã©ãŒãããã®ONNX 3 ããã®å€æããPyTorchã¢ãã«ããã®çŽæ¥å€æãå¯èœã§ãã å®ã¯NVIDIAãå
¬åŒããã°ã§PARSeqãTensorRTåããèšäºãå
¬éããŠãã 4 ã®ã§ããã PARSeqããã®äŸåå
ã®PyTorchã®ããŒãžã§ã³ãå€ããã®ãŸãŸã§ã¯åäœããªããããæ¬çš¿ã§ã¯ææ°ç(PyTorch 2.10, PARSeq 2024幎2æç)ã䜿ã£ãTensorRTåã®æµãã玹ä»ããŸãã PyTorch Lightningã«ããã¢ãã«å€æ PARSeqã¯PyTorchã«ãã£ãŠå®è£
ãããã¢ãã«ãPyTorch-Lightningã§å¶åŸ¡ããŠãããONNXãTensorRTãžã®å€æã¯PyTorch-LightningãæäŸãã颿°ãå©çšã§ããŸãã NVIDIAã®ããã°ã§ã to_onnx() ãå©çšããŠäžåºŠONNXåããã®ã¡ãtrtexecãšåŒã°ããããŒã«ã䜿ã£ãŠONNXããTensorRTãžå€æããŠããŸãã ä»å㯠to_tensorrt() ãå©çšããŠãã¢ãã«ãçŽæ¥TensorRTã«å€æããŠã¿ãŸãã import torch parseq = torch.hub.load( 'baudm/parseq' , 'parseq' , pretrained= True ).eval() parseq.model.refine_iters = 0 # Iterative refinementãç¡å¹å parseq.model.decode_ar = False # Non-autoregressive mode output_path = "engine.pt2" img = torch.randn( 1 , 3 , 32 , 128 ) parseq.to_tensorrt(output_path, img, ir= "dynamo" ) ããã§ç¡äºTensorRTã¢ãã« engine.pt2 ã«å€æã§ããŸããããã®ã¢ãã«ã¯ä»¥äžã®ããã«åŒã³åºãããšãã§ããŸãã import torch import torch_tensorrt # <- å¿
é parseq = torch.export.load( "engine.pt2" ).module() img = torch.randn( 1 , 3 , 32 , 128 ).cuda() parseq(img) # torch.Size([1, 26, 95]) 26ã¯äžåºŠã«æšæž¬å¯èœãªæåæ°ã95ã¯å¯Ÿå¿æåçš® AutoregressiveãšIterative refinementãTensorRTåã§ããªãåé¡ ããããªããããã®æ¹æ³ã§ã¯Autoregressive( decode_ar=True )ãŸãã¯Iterative refinement( refine_iters>0 )ã«å¯Ÿå¿ããã¢ãã«ãäœãããšãããšãšã©ãŒã«ãªã£ãŠããŸããŸãã è«æã§ã¯Non-autoregressiveããAutoregressiveã®æ¹ãé«ç²ŸåºŠ 5 ãšãããŠããããŸãIterative refinementã1åé©çšããã ãã§ãããªãã«ç²ŸåºŠãåäžããããããã²ãããã®ã¢ãŒããTensorRTã§æŽ»çšãããã§ãã ããã§PARSeqã®å®è£
ãæ¹é ãTensorRTåã«ææŠããŸããã Autoregressive modeã®TensorRTå ãŸãå
ã»ã©ãšåãæ¹æ³ã§ã¯ã©ãã§èœã¡ãããã¿ãŠã¿ãŸãã import torch parseq = torch.hub.load( 'baudm/parseq' , 'parseq' , pretrained= True ).eval() parseq.model.refine_iters = 0 parseq.model.decode_ar = True # AR mode output_path = "engine.pt2" img = torch.randn( 1 , 3 , 32 , 128 ) parseq.to_tensorrt(output_path, img, ir= "dynamo" ) 衚瀺ããããšã©ãŒã¯ä»¥äžã®ãšããã§ãã File "/root/.cache/torch/hub/baudm_parseq_main/strhub/models/parseq/model.py", line 144, in forward if testing and (tgt_in == tokenizer.eos_id).any(dim=-1).all(): ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ... torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not guard on data-dependent expression Eq(u0, 1) (unhinted: Eq(u0, 1)). (Size-like symbols: none) ããã¯ãæç« ã®çµäºã瀺ãEOSããŒã¯ã³ãåºããçæã忢ãããåŠçã®éšåã§ãããã©ããifæã«ããåå²ã¯TensorRTãšçžæ§ãæªãããã§ãã ãããããã¯ïŒæåçæãç¹°ãè¿ãAutoregressive modeã§ã¯å¿
é ã®åŠçã§ããããã1æåçæããå®è£
ã®ã¿ãTensorRTåããç¹°ãè¿ãéšåã¯ã¢ãã«ã®å€åŽã§ããããã«å€ããŠã¿ãŸãã import pytorch_lightning as pl from torch import Tensor from typing import Optional class PARSeqEncoder (pl.LightningModule): def __init__ (self, model): super ().__init__() self.encoder = model.encoder def forward (self, images: Tensor) -> Tensor: memory = self.encoder(images) return memory class PARSeqDecoder (pl.LightningModule): def __init__ (self, tokenizer, model): super ().__init__() self.tokenizer = tokenizer self.max_label_length = model.max_label_length self.text_embed = model.text_embed self.pos_queries = model.pos_queries self.decoder = model.decoder self.head = model.head def forward (self, memory: Tensor, input_ids: Tensor) -> Tensor: B, S = input_ids.size( 0 ), input_ids.size( 1 ) null_ctx = self.text_embed(input_ids[:, : 1 ]) tgt_emb = self.pos_queries[:, :S- 1 ] + self.text_embed(input_ids[:, 1 :]) tgt_emb = torch.cat([null_ctx, tgt_emb], dim= 1 ) tgt_query = self.pos_queries[:, S- 1 :S].expand(B, - 1 , - 1 ) tgt_mask = torch.triu(torch.ones((S, S), dtype=torch.bool), 1 ).to(tgt_emb.device) decoder_outputs = self.decoder(tgt_query, tgt_emb, memory, content_mask=tgt_mask) return self.head(decoder_outputs) ããã§ãšã³ã³ãŒããšãã³ãŒããåãé¢ãããŠããŸããããã¯ãšã³ã³ãŒããäžåºŠå®è¡ããã®ã¡ããã³ãŒããEOSããŒã¯ã³ãåºããŸã§ç¹°ãè¿ãå¿
èŠãããããã§ãã æšè«ã¯ä»¥äžã®ããã«ãªããŸãã img_transform = T.Compose([ T.Resize(( 32 , 128 ), T.InterpolationMode.BICUBIC), T.ToTensor(), T.Normalize( 0.5 , 0.5 ), ]) _parseq = torch.hub.load( 'baudm/parseq' , 'parseq' , pretrained= True ).eval() bos_id = _parseq.tokenizer.bos_id pad_id = _parseq.tokenizer.pad_id eos_id = _parseq.tokenizer.eos_id parseq_encoder = PARSeqEncoder(_parseq.model) parseq_decoder = PARSeqDecoder(_parseq.tokenizer, _parseq.model) img = Image.open( "world.png" ).convert( "RGB" ) img = img_transform(img).unsqueeze( 0 ) with torch.no_grad(): num_steps = _parseq.model.max_label_length + 1 input_ids = torch.full(( 1 , num_steps), pad_id, dtype=torch.long) input_ids[:, 0 ] = bos_id memory = parseq_encoder(img) preds = [] for i in range (num_steps- 1 ): j = i + 1 logit = parseq_decoder(memory, input_ids[:, :j]) preds.append(logit.softmax(- 1 )) input_ids[:, j:j+ 1 ] = logit.argmax(- 1 ) if (input_ids == eos_id).any(dim=- 1 ).all(): break label, confidence = _parseq.tokenizer.decode(torch.cat(preds, dim= 1 )) print (f "AR result: {label[0]}" ) ãããŠå€æã¯æ¬¡ã®ããã«è¡ããŸãã input_ids ã®é·ãã¯äŒžã³çž®ã¿ããããæçã»æé·ãæå®ããŠããå¿
èŠããããŸãã parseq_encoder.to_tensorrt( "encoder.pt2" , img, ir= "dynamo" ) decoder_input_ids = torch_tensorrt.Input( min_shape=[ 1 , 1 ], opt_shape=[ 1 , num_steps], max_shape=[ 1 , num_steps], dtype=torch.int64) encoder_outputs = torch_tensorrt.Input( min_shape=[ 1 , 128 , 384 ], opt_shape=[ 1 , 128 , 384 ], max_shape=[ 1 , 128 , 384 ], dtype=torch.float32) parseq_decoder.to_tensorrt( "decoder.pt2" , (encoder_outputs, decoder_input_ids), ir= "dynamo" ) TorchDynamoã®æ©å«ããšã ããããªããããªããããã¯ãã³ãŒã( PARSeqDecoder )ã®å€æã«å€±æããŸããæ¬æ¥å
¥åãã input_ids ã®ããŒã¯ã³é·ã¯1以äžããã°åäœããã¯ãã§ããã以äžã®ããã«3以äžã«éå®ããªãããšãããšã©ãŒãåºãŠããŸãã - Not all values of _1 = L['input_ids'].size()[1] in the specified range _1 <= 26 satisfy the generated guard 3 <= L['input_ids'].size()[1] and L['input_ids'].size()[1] <= 26 Suggested fixes: _1 = Dim('_1', min=3, max=26) ããã¯TensorRTåãããåã®ãTorchDynamoããœãŒã¹ã³ãŒããè§£æãããšãã«çºçããŠãããšã©ãŒãªã®ã§ãããã©ããåå ãªã®ããTorchDynamoã䜿ã£ãŠæ¢ã£ãŠã¿ãŸãã from torch_tensorrt.dynamo.utils import get_torch_inputs, to_torch_device from torch_tensorrt.dynamo._tracer import get_dynamic_shapes_args from torch.export import Dim, export, draft_export arg_inputs = (encoder_outputs, decoder_input_ids) parseq_decoder.to( "cuda" ) device = to_torch_device( "cuda" ) torch_arg_inputs = get_torch_inputs(arg_inputs, device) dynamic_shapes = get_dynamic_shapes_args(parseq_decoder, arg_inputs) ep = draft_export( # ãšã©ãŒãèµ·ããŠãæåŸãŸã§è§£æãããããšã§å
šãŠã®ãšã©ãŒãåéãã parseq_decoder, tuple (torch_arg_inputs), dynamic_shapes=dynamic_shapes, ) print (ep._report) ãããšä»¥äžã®ãããªèŠåã確èªã§ããŸãã ################################################################################################### WARNING: 2 issue(s) found during export, and it was not able to soundly produce a graph. Please follow the instructions to fix the errors. ################################################################################################### 1. Guard Added. A guard was added during tracing, which might've resulted in some incorrect tracing or constraint violation error. Specifically, this guard was added: Ne(s70 - 1, 1), where {'s70': "L['input_ids'].size()[1]"}. This occurred at the following stacktrace: File /opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py, lineno 1776, in _wrapped_call_impl File /opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py, lineno 1787, in _call_impl File /workspace/src/ar_deploy_decoder.py, lineno 31, in forward tgt_emb = self.pos_queries[:, :S-1] + self.text_embed(input_ids[:, 1:]): Locals: self: [None] S: ['s70'] input_ids: ['Tensor(shape: torch.Size([1, s70]), stride: (s70, 1), storage_offset: 0)'] Symbols: s70: L['input_ids'].size()[1] And the following framework stacktrace: File /opt/venv/lib/python3.12/site-packages/torch/_prims_common/__init__.py, lineno 404, in is_contiguous_for_memory_format File /opt/venv/lib/python3.12/site-packages/torch/_prims_common/__init__.py, lineno 317, in is_contiguous File /opt/venv/lib/python3.12/site-packages/torch/_prims_common/__init__.py, lineno 277, in check_contiguous_sizes_strides if maybe_guard_or_false(x == 1): (以äžçç¥) ãã³ãœã«ã S-1 ã®é·ãã«ã¹ã©ã€ã¹ãããšããã§ S-1 != 1 ãšããå¶çŽãDynamoã«ãã£ãŠå°å
¥ãããŠããŸãã ã©ãããã¹ã©ã€ã¹ãããæé·ã1ã« ãªããã å¯å€é·ãã³ãœã«ã¯åé¡ãããããã§ããïŒãããã0/1-specialization 6 ãšåŒã°ããåŠçãšé¢ä¿ãããã®ã§ããããªããããªã£ãŠããã®ãã¯ããåãããŸãã...ïŒ ããã§ã¹ã©ã€ã¹ãè¡ããªã圢ã«å®è£
ãçŽããŠãããŸãã class PARSeqDecoder (pl.LightningModule): def __init__ (self, tokenizer, model): super ().__init__() self.tokenizer = tokenizer self.max_label_length = model.max_label_length self.text_embed = model.text_embed # self.pos_queries = model.pos_queries self.prefixed_pos_queries = torch.nn.Parameter(torch.cat([torch.zeros_like(model.pos_queries)[:,: 1 ], model.pos_queries], dim= 1 )) self.decoder = model.decoder self.head = model.head def forward (self, memory: Tensor, input_ids: Tensor) -> Tensor: B, S = input_ids.size( 0 ), input_ids.size( 1 ) tgt_emb = self.prefixed_pos_queries[:, :S] + self.text_embed(input_ids) tgt_query = self.prefixed_pos_queries[:, S:S+ 1 ].expand(B, - 1 , - 1 ) tgt_mask = torch.triu(torch.ones((S, S), dtype=torch.bool), 1 ).to(tgt_emb.device) decoder_outputs = self.decoder(tgt_query, tgt_emb, memory, content_mask=tgt_mask) return self.head(decoder_outputs) ããã§ç¡äºå€æãéãããã«ãªããŸããã Iterative refinementã®TensorRTå æ¬¡ã«Iterative refinementãè¡ããã³ãŒãã®TensorRTåãè¡ããŸãã å
ã®PARSeqå®è£
ããrefinementãè¡ãç®æãåãåºãPyTorch Lightningã§ã©ããããŸãã class PARSeqRefiner (pl.LightningModule): def __init__ (self, tokenizer, model): super ().__init__() self.tokenizer = tokenizer self.max_label_length = model.max_label_length self.text_embed = model.text_embed self.prefixed_pos_queries = torch.nn.Parameter(torch.cat([torch.zeros_like(model.pos_queries)[:,: 1 ], model.pos_queries], dim= 1 )) self.pos_queries = model.pos_queries self.decoder = model.decoder self.head = model.head def forward (self, memory: Tensor, input_ids: Tensor) -> Tensor: B, S = input_ids.size( 0 ), input_ids.size( 1 ) tgt_emb = self.prefixed_pos_queries[:, :S] + self.text_embed(input_ids) tgt_query = self.pos_queries tgt_mask = torch.triu(torch.ones((S, S), dtype=torch.bool), 1 ).to(tgt_emb.device) tgt_mask[torch.triu(torch.ones((S, S), dtype=torch.bool, device=tgt_emb.device), 2 )] = 0 tgt_padding_mask = (input_ids == self.tokenizer.eos_id).int().cumsum(- 1 ) > 0 decoder_outputs = self.decoder(tgt_query, tgt_emb, memory, query_mask=tgt_mask, content_mask=tgt_mask, content_key_padding_mask=tgt_padding_mask) return self.head(decoder_outputs) refiner_input_ids = torch_tensorrt.Input( min_shape=[ 1 , num_steps], opt_shape=[ 1 , num_steps], max_shape=[ 1 , num_steps], dtype=torch.int64) print ( "==== export refiner ====" ) parseq_refiner.to_tensorrt( "refiner.pt2" , (encoder_outputs, refiner_input_ids), ir= "dynamo" ) ãã¡ãã¯å
¥åããŒã¯ã³ã䌞ã³çž®ã¿ããªãã®ãããçŽ çŽã«å€æã§ããŸããã è©äŸ¡ æåŸã«TensorRTåã«ãã£ãŠã©ããããéããªã£ãããã¿ãŠã¿ãŸãã OCRã®ãã³ãããŒã¯ã§ããIIIT-5Kã«å¯ŸããŠããŸããŸãªèšå®ã§æšè«ãã1æãããã®ã¬ã€ãã³ã·ãH200 GPU 1å°ã§èšæž¬ããŸããã çµæã¯æ¬¡ã®å³ã®ããã«ãªããŸããã äŸãã°Autoregressive(AR)ã¢ãŒãã»iterative refinementç¡ãã§ã¯TensorRT倿ã«ãã£ãŠ2.58åã®é«éåã Non-Autoregressive(NAR)ã¢ãŒãã§ã¯3.07åã®é«éåãéæããŸããã ã°ã©ãã®åŸãããiterative refinementã軜éã«ãªã£ãŠããããšãåãããŸãã ãŸãšã ä»åã®å®éšã§ã¯è»œéã§é«æ§èœãªOCRã¢ãã«ã§ããPARSeqãææ°ã®ç°å¢ã§TensorRTåããŠã¿ãŸããã ãã®éãæç« çæãªã©ã§ããçšãããããã³ãŒãã¯å
¥åãµã€ãºãåçã«å€åãããã倿ã«äžçãããã©ã€ãã©ãªãåŠçãããããããªããã°ã©ã ã«æžãæããå¿
èŠãããããšã玹ä»ããŸããã https://github.com/baudm/parseq ↩ https://developer.nvidia.com/tensorrt ↩ https://onnx.ai ↩ https://developer.nvidia.com/blog/robust-scene-text-detection-and-recognition-inference-optimization/ ↩ https://arxiv.org/abs/2207.06966 Appendix H ↩ https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/torch.compiler_dynamo_deepdive.html#are-always-specialized ↩