| 88 | return None |
| 89 | |
| 90 | def getT5DecodingPlugin(arg): |
| 91 | nBatchSize = arg['batch_size'] |
| 92 | nMaxSeqLen = arg['max_seq_len'] |
| 93 | nMemMaxSeqLen = arg['max_seq_len'] |
| 94 | nBeamSize = arg['beam_width'] |
| 95 | useFP16 = int(arg['data_type'] == 'fp16') |
| 96 | ckpt_path = arg['ckpt_path'].encode() |
| 97 | for c in trt.get_plugin_registry().plugin_creator_list: |
| 98 | if c.name == 'T5DecodingPlugin': |
| 99 | pList = [ |
| 100 | trt.PluginField('max_batch_size', np.int32(nBatchSize), npToPFT[np.int32]), |
| 101 | trt.PluginField('max_seq_len', np.int32(nMaxSeqLen), npToPFT[np.int32]), |
| 102 | trt.PluginField('mem_max_seq_len', np.int32(nMaxSeqLen), npToPFT[np.int32]), |
| 103 | trt.PluginField('beam_width', np.int32(nBeamSize), npToPFT[np.int32]), |
| 104 | trt.PluginField('useFP16', np.int32(useFP16), npToPFT[np.int32]), |
| 105 | trt.PluginField('ckpt_path', ckpt_path, trt.PluginFieldType.CHAR), |
| 106 | ] |
| 107 | return c.create_plugin(c.name, trt.PluginFieldCollection(pList)) |
| 108 | return None |
| 109 | |
| 110 | def buildEngine(logger, arg, trtFileName): |
| 111 | builder = trt.Builder(logger) |