torch2trt_dynamic: No warnings, but [TensorRT] ERROR: INVALID_ARGUMENT: Cannot find binding of given name: input_0

Actually, Nope, like most tracing based deploy tools, this project does not support python control flow statements such as if or while etc. That means even you can convert torch.Tensor.__bool__, the if statement still won’t work. Try:

output = input * alpha + other * ( 1-alpha)

to avoid control flow.

Thanks for your reply! I’ve fixed all the bugs in my code. I’d like to summarize two points:

1. The ‘cannot find input’ bug is always caused by some unsupported layer in your model. You have to find it and write the corresponding op if necessary.
2. One problem is that if you use the from torch2trt_dynamic import torch2trt_dynamic in your python code, you will not see the Traceback, you will only see one line: segmentation fault (core dumped). This bothered me so much during the debug process. I will give my GRU converter here:

import tensorrt as trt
import torch
from torch import nn
from torch2trt_dynamic.torch2trt_dynamic import *


@tensorrt_converter('torch.nn.GRU.forward')
def convert_gru(ctx):

    module = ctx.method_args[0]
    input_tensor = ctx.method_args[1]
    if len(ctx.method_args) == 3:
        init_state_tensor = ctx.method_args[2]
    output_0, output_1 = ctx.method_return[0], ctx.method_return[1]

    layer_count = module.num_layers
    hidden_size = module.hidden_size
    batch_first = module.batch_first
    max_seq_length = input_tensor.shape[1] if batch_first else input_tensor.shape[0]
    op = trt.RNNOperation.GRU
    layer = ctx.network.add_rnn_v2(input_tensor._trt, layer_count, hidden_size, max_seq_length, op)
    if len(ctx.method_args) == 3:
        layer.hidden_state = init_state_tensor._trt

    if module.bidirectional is True:
        layer.direction = trt.RNNDirection.BIDIRECTION
    for i in range(layer_count):
        iw = getattr(module, "weight_ih_l%s" % i).detach().cpu().numpy()
        hw = getattr(module, "weight_hh_l%s" % i).detach().cpu().numpy()

        rela_index = 2 * i if module.bidirectional is True else i

        layer.set_weights_for_gate(rela_index, trt.RNNGateType.UPDATE, True, iw[:hidden_size, :].copy())
        layer.set_weights_for_gate(rela_index, trt.RNNGateType.RESET, True, iw[hidden_size: hidden_size * 2, :].copy())
        layer.set_weights_for_gate(rela_index, trt.RNNGateType.HIDDEN, True, iw[hidden_size * 2: hidden_size * 3, :].copy())

        layer.set_weights_for_gate(rela_index, trt.RNNGateType.UPDATE, False, hw[:hidden_size, :].copy())
        layer.set_weights_for_gate(rela_index, trt.RNNGateType.RESET, False, hw[hidden_size: hidden_size * 2, :].copy())
        layer.set_weights_for_gate(rela_index, trt.RNNGateType.HIDDEN, False, hw[hidden_size * 2: hidden_size * 3, :].copy())

        ib = getattr(module, "bias_ih_l%s" % i).detach().cpu().numpy()
        hb = getattr(module, "bias_hh_l%s" % i).detach().cpu().numpy()

        layer.set_bias_for_gate(rela_index, trt.RNNGateType.UPDATE, True, ib[:hidden_size].copy())
        layer.set_bias_for_gate(rela_index, trt.RNNGateType.RESET, True, ib[hidden_size:hidden_size * 2].copy())
        layer.set_bias_for_gate(rela_index, trt.RNNGateType.HIDDEN, True, ib[hidden_size * 2: hidden_size * 3].copy())

        layer.set_bias_for_gate(rela_index, trt.RNNGateType.UPDATE, False, hb[:hidden_size].copy())
        layer.set_bias_for_gate(rela_index, trt.RNNGateType.RESET, False, hb[hidden_size:hidden_size * 2].copy())
        layer.set_bias_for_gate(rela_index, trt.RNNGateType.HIDDEN, False, hb[hidden_size * 2: hidden_size * 3].copy())

        if module.bidirectional is True:
            # ================reverse=====================
            iw_r = getattr(module, "weight_ih_l%s_reverse" % i).detach().cpu().numpy()
            hw_r = getattr(module, "weight_hh_l%s_reverse" % i).detach().cpu().numpy()

            layer.set_weights_for_gate(2 * i + 1, trt.RNNGateType.UPDATE, True, iw_r[:hidden_size, :].copy())
            layer.set_weights_for_gate(2 * i + 1, trt.RNNGateType.RESET, True, iw_r[hidden_size:hidden_size * 2, :].copy())
            layer.set_weights_for_gate(2 * i + 1, trt.RNNGateType.HIDDEN, True, iw_r[hidden_size * 2: hidden_size * 3, :].copy())

            layer.set_weights_for_gate(2 * i + 1, trt.RNNGateType.UPDATE, False, hw_r[:hidden_size, :].copy())
            layer.set_weights_for_gate(2 * i + 1, trt.RNNGateType.RESET, False, hw_r[hidden_size:hidden_size * 2, :].copy())
            layer.set_weights_for_gate(2 * i + 1, trt.RNNGateType.HIDDEN, False, hw_r[hidden_size * 2: hidden_size * 3, :].copy())

            ib_r = getattr(module, "bias_ih_l%s_reverse" % i).detach().cpu().numpy()
            hb_r = getattr(module, "bias_hh_l%s_reverse" % i).detach().cpu().numpy()

            layer.set_bias_for_gate(2 * i + 1, trt.RNNGateType.UPDATE, True, ib_r[:hidden_size].copy())
            layer.set_bias_for_gate(2 * i + 1, trt.RNNGateType.RESET, True, ib_r[hidden_size:hidden_size * 2].copy())
            layer.set_bias_for_gate(2 * i + 1, trt.RNNGateType.HIDDEN, True, ib_r[hidden_size * 2: hidden_size * 3].copy())

            layer.set_bias_for_gate(2 * i + 1, trt.RNNGateType.UPDATE, False, hb_r[:hidden_size].copy())
            layer.set_bias_for_gate(2 * i + 1, trt.RNNGateType.RESET, False, hb_r[hidden_size:hidden_size * 2].copy())
            layer.set_bias_for_gate(2 * i + 1, trt.RNNGateType.HIDDEN, False, hb_r[hidden_size * 2: hidden_size * 3].copy())

    gru_output_0 = layer.get_output(0)
    gru_output_1 = layer.get_output(1)
    output_0._trt = gru_output_0
    output_1._trt = gru_output_1


def main():
    class TestNet(torch.nn.Module):
        def __init__(self):
            super(TestNet, self).__init__()
            self.gru = nn.GRU(10, 20, 1, batch_first=True, bidirectional=False)

        def forward(self, x, init):
            out, out1 = self.gru(x, init)
            return out, out1

    net = TestNet()
    net = net.cuda()

    x = torch.randn(1, 5, 10).cuda()
    init = torch.randn(1, 1, 20).cuda()
    output_0, hn_0 = net(x, init)

    trt_net = torch2trt_dynamic(net, [x, init],  max_workspace_size=1 << 28)

    output_1, hn_1 = trt_net(x, init)

    print(torch.max(torch.abs(output_0 - output_1)))
    print(torch.max(torch.abs(hn_0 - hn_1)))


if __name__ == "__main__":
    main()

Jetson Nano, Pytorch 1.8.0, TensorRT 7.1.3.0 Thanks to https://github.com/NVIDIA-AI-IOT/torch2trt/issues/144#issuecomment-553321172