problem about function : view()

2:torch.Size([4, 256]) (256,) t–>trt: (256,) [TensorRT] ERROR: (Unnamed Layer* 177) [Shuffle]: uninferred dimensions are not an exact divisor of input dimensions, so inferred dimension cannot be calculated torch.Size([1, 4, 256]) (0) t–>trt: ()

when I run view() from tensor[4,256] to [1,4,256], some problem happens----->t._trt.shape becomes 0!!!

code: in torch2trt.py `def add_missing_trt_tensors(network, tensors): “”“Creates missing TensorRT tensors as constants and attaches them to the Torch Tensors”“” trt_tensors = [None] * len(tensors)

dtype = check_torch_dtype(*tensors)

for i, t in enumerate(tensors):
    trt_tensor = None

    # GET TRT TENSOR (OR CREATE TRT CONSTANT)

    # get tensor w/ _trt
    # or... add constant for scalar primitive
    if isinstance(t, float) or isinstance(t, int):
        shape = (1,)
        scalar = t * torch.ones(shape, dtype=dtype).cpu().numpy()
        trt_tensor = network.add_constant(shape, scalar).get_output(0)
    elif hasattr(t, "_trt"):
        print('2:',t.shape,t._trt.shape)
        trt_tensor = t._trt
        print('t-->trt:',trt_tensor.shape)

    # or... add constant for leaf tensor w/o _trt
    else:
        # remove all preceding ones, these can be re-inserted later when broadcasting
        num_preceding_ones = 0
        for j in range(t.ndim):
            if int(t.shape[j]) == 1:
                num_preceding_ones += 1
            else:
                break
        shape = tuple(t.shape[num_preceding_ones:])

        weight = t.detach().cpu().numpy()
        t._trt = network.add_constant(shape, weight).get_output(0)
        trt_tensor = t._trt


    assert trt_tensor is not None
    trt_tensors[i] = trt_tensor
return trt_tensors`

[jaybdub]

Hmm. I see, so in you’re model you’re actually running on a batch of 4 images, and then permuting the batch dimension.

Unfortunately, permuting the batch dimension is not currently supported by torch2trt.

That said, if resnet50 is your backbone on a batch of 4 images, I would expect this final layer is a small fraction of the computation. You could convert the backbone model with torch2trt, and run the linear layer in PyTorch. For example,

class MyModel(nn.Module):
    def __init__(self, ...):
        self.backbone = resnet50(...)
        self.linear = nn.Linear(...)

    def forward(self, x):
        x = self.backbone(x)
        x = x.view(1, 1024)
        x = self.linear(x)
        return x

and then convert like

model = MyModel()
backbone_trt = torch2trt(model.backbone, [data], ...)
model.backbone = backbone_trt

torch.save(model.state_dict(), 'model_trt.pth')

which could be later loaded like

model = MyModel()
model.backbone = TRTModule()
model.load_state_dict(torch.load('model_trt.pth'))

You could also convert the linear layer with torch2trt, but you may not see much benefit.

Please let me know if this helps or you run into any issues.

Best, John

[lxy443626128]

@jaybdub It works as you suggested, thanks!!!