Can't use custom plugin in trt 7 loop

I implement a custom gather plugin (IPluginV2Ext) (equivalent to tensorflow’s gather), but get error when use it inside a loop:
[TensorRT] ERROR: …/builder/myelin/codeGenerator.cpp (338) - Myelin Error in compileGraph: 8 (myelinUninitializedTensor : Operation “(Unnamed Layer_ 38) [Shuffle]first_transpose" uses tensor "(Unnamed Layer 37) [PluginV2Ext]output_0’" which has not been defined.
[TensorRT] ERROR: …/builder/myelin/codeGenerator.cpp (338) - Myelin Error in compileGraph: 8 (myelinUninitializedTensor : Operation "(Unnamed Layer
38) [Shuffle]first_transpose" uses tensor "(Unnamed Layer 37) [PluginV2Ext]_output_0’” which has not been defined.
code snippet:

def _gather(x, indices, axis, batch_dims=0, name=None):
    net = get_trt_network()
    config = {
        "axis": axis,
        "batch_dims": batch_dims,
    plugin = create_plugin("GatherCustom", {"config": json.dumps(config)})
    layer = net.add_plugin_v2([x, indices], plugin)
    output = layer.get_output(0)
    if name is not None: = name = name
    return output

def gather(x, indices, axis, batch_dims=0, name=None):
    return Tensor(_gather(to_wrapped(x), to_wrapped(indices), axis, batch_dims, name))

def get_loop(inp: T.Tensor, h: T.Tensor, c: T.Tensor, feature: T.Tensor, params):
    # feature: [1, 66, 9216]
    # inp: [1, 66]
    net = T.get_trt_network()
    loop = net.add_loop()
        T.tensor(np.array(41, dtype=np.int32)).wrapped, trt.TripLimit.COUNT)
    rec_inp = loop.add_recurrence(inp.wrapped)
    rec_h = loop.add_recurrence(h.wrapped)
    rec_c = loop.add_recurrence(c.wrapped)
    inp_rec = rec_inp.get_output(0)
    h_rec = rec_h.get_output(0)
    c_rec = rec_c.get_output(0)

    output, [h_next, c_next] = get_lstm_unit(inp_rec, [h_rec, c_rec], params)
    # output: [1, 1]
    output = T.maximum(output, feature.shape[2] - 1)
    output = T.minimum(output, 0)
    input_data = T.gather(feature, output, axis=2)
    # input_data = T.gather(feature, T.tensor(np.array([[1]], dtype=np.int32)), axis=2)
    input_data = input_data.transpose(0, 2, 1).reshape(inp.shape)
    rec_inp.set_input(1, input_data.wrapped)
    rec_h.set_input(1, h_next.wrapped)
    rec_c.set_input(1, c_next.wrapped)

    inp_out = loop.add_loop_output(inp_rec, trt.LoopOutput.LAST_VALUE)
    h_out = loop.add_loop_output(h_rec, trt.LoopOutput.LAST_VALUE)
    c_out = loop.add_loop_output(c_rec, trt.LoopOutput.LAST_VALUE)

    return T.Tensor(inp_out.get_output(0)), T.Tensor(h_out.get_output(0)), T.Tensor(c_out.get_output(0))

The “T.gather” is plugin layer, other layers are native. If I use constant indices instead of loop recurrence variable, the code works fine.


Currently there’s no support for using user-defined plugin inside a loop in TRT 7.



I have the same problem. I need to implement some custom computation through plugin and put it in a big loop at a 100,000 level. It seems I need to put the loop outside the network graph and run enqueue() 100,000 times, Right ? Will that affect performance heavily ? And any plan to support custom layer in ILoop ?


Hi SuniJB,

I’m having the same problem. I am now using TensorRT- and according the interior of the loop can have IPluginV2Layer. As hl1008 mentioned, I’m afraid it’ll be inefficient to use repeated enqueue call instead of a loop.

Thank you in advance.