PyTorch model converting to TensorRT issue

I have trained model. Here is the class

class UNet(nn.Module):
    def __init__(self, n_channels=1, n_classes=2):
        super(UNet, self).__init__()
        self.inc = down(1, 32)
        self.down1 = down(32, 64)
        self.down2 = down(64, 128)
        self.down3 = down(128, 256)
        self.down4 = down(256, 512)
        self.up1 = up(512, 512)
        self.up2 = up(512 + 256, 256)
        self.up3 = up(256 + 128, 128)
        self.up4 = up(128 + 64, 64)
        self.outc = outconv(64 + 32, 1)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        x = self.outc(x)
        return torch.tanh(x)

class down(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(down, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, stride=2, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU())
        
    def forward(self, x):
        x = self.conv(x)
        return x


class up(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(up, self).__init__()
        self.conv = nn.Sequential(
            nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2, padding=0),
            nn.BatchNorm2d(out_ch),
            nn.Dropout(p=0.5),
            nn.ReLU())

    def forward(self, x1, x2):
        x = self.conv(x1)
        x = torch.cat([x2, x], dim=1)
        return x


class outconv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(outconv, self).__init__()
        self.conv = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2, padding=0)

    def forward(self, x):
        x = self.conv(x)
        return x

I try to convert it to tensorRt with torch2trt converter. Here is the log of process from jupyter notebook console

[TensorRT] INFO: Applying generic optimizations to the graph for inference.
[TensorRT] INFO: Original: 33 layers
[TensorRT] INFO: After dead-layer removal: 33 layers
[TensorRT] INFO: Fusing convolution weights from (Unnamed Layer* 0) [Convolution] with scale (Unnamed Layer* 1) [Scale]
[TensorRT] INFO: Fusing convolution weights from (Unnamed Layer* 3) [Convolution] with scale (Unnamed Layer* 4) [Scale]
[TensorRT] INFO: Fusing convolution weights from (Unnamed Layer* 6) [Convolution] with scale (Unnamed Layer* 7) [Scale]
[TensorRT] INFO: Fusing convolution weights from (Unnamed Layer* 9) [Convolution] with scale (Unnamed Layer* 10) [Scale]
[TensorRT] INFO: Fusing convolution weights from (Unnamed Layer* 12) [Convolution] with scale (Unnamed Layer* 13) [Scale]
[TensorRT] INFO: After scale fusion: 28 layers
[TensorRT] INFO: Fusing (Unnamed Layer* 0) [Convolution] with (Unnamed Layer* 2) [Activation]
[TensorRT] INFO: Fusing (Unnamed Layer* 3) [Convolution] with (Unnamed Layer* 5) [Activation]
[TensorRT] INFO: Fusing (Unnamed Layer* 6) [Convolution] with (Unnamed Layer* 8) [Activation]
[TensorRT] INFO: Fusing (Unnamed Layer* 9) [Convolution] with (Unnamed Layer* 11) [Activation]
[TensorRT] INFO: Fusing (Unnamed Layer* 12) [Convolution] with (Unnamed Layer* 14) [Activation]
[TensorRT] INFO: Fusing (Unnamed Layer* 16) [Scale] with (Unnamed Layer* 17) [Activation]
[TensorRT] INFO: Fusing (Unnamed Layer* 20) [Scale] with (Unnamed Layer* 21) [Activation]
[TensorRT] INFO: Fusing (Unnamed Layer* 24) [Scale] with (Unnamed Layer* 25) [Activation]
[TensorRT] INFO: Fusing (Unnamed Layer* 28) [Scale] with (Unnamed Layer* 29) [Activation]
[TensorRT] INFO: After vertical fusions: 19 layers
[TensorRT] INFO: After swap: 19 layers
[TensorRT] INFO: After final dead-layer removal: 19 layers
[TensorRT] INFO: After tensor merging: 19 layers
[TensorRT] INFO: Eliminating concatenation (Unnamed Layer* 18) [Concatenation]
[TensorRT] INFO: Retargeting (Unnamed Layer* 11) [Activation]_output to (Unnamed Layer* 18) [Concatenation]_output
[TensorRT] INFO: Retargeting (Unnamed Layer* 17) [Activation]_output to (Unnamed Layer* 18) [Concatenation]_output
[TensorRT] INFO: Eliminating concatenation (Unnamed Layer* 22) [Concatenation]
[TensorRT] INFO: Retargeting (Unnamed Layer* 8) [Activation]_output to (Unnamed Layer* 22) [Concatenation]_output
[TensorRT] INFO: Retargeting (Unnamed Layer* 21) [Activation]_output to (Unnamed Layer* 22) [Concatenation]_output
[TensorRT] INFO: Eliminating concatenation (Unnamed Layer* 26) [Concatenation]
[TensorRT] INFO: Retargeting (Unnamed Layer* 5) [Activation]_output to (Unnamed Layer* 26) [Concatenation]_output
[TensorRT] INFO: Retargeting (Unnamed Layer* 25) [Activation]_output to (Unnamed Layer* 26) [Concatenation]_output
[TensorRT] INFO: Eliminating concatenation (Unnamed Layer* 30) [Concatenation]
[TensorRT] INFO: Retargeting (Unnamed Layer* 2) [Activation]_output to (Unnamed Layer* 30) [Concatenation]_output
[TensorRT] INFO: Retargeting (Unnamed Layer* 29) [Activation]_output to (Unnamed Layer* 30) [Concatenation]_output
[TensorRT] INFO: After concat removal: 15 layers
[TensorRT] INFO: Graph construction and optimization completed in 0.0351903 seconds.
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 0) [Convolution] + (Unnamed Layer* 2) [Activation](3)
[TensorRT] INFO: Tactic 0 time 0.54672
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 0) [Convolution] + (Unnamed Layer* 2) [Activation](2)
[TensorRT] INFO: Tactic 6 time 2.80048
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 0) [Convolution] + (Unnamed Layer* 2) [Activation](14)
[TensorRT] INFO: Tactic 3146172331490511787 time 2.18669
[TensorRT] INFO: Tactic 3528302785056538033 time 0.72176
[TensorRT] INFO: Tactic 5443600094180187792 time 0.6728
[TensorRT] INFO: Tactic 5824828673459742858 time 4.45536
[TensorRT] INFO: Tactic -6618588952828687390 time 0.69616
[TensorRT] INFO: Tactic -6362554771847758902 time 1.11968
[TensorRT] INFO: Tactic -2701242286872672544 time 2.17491
[TensorRT] INFO: Tactic -2535759802710599445 time 1.08195
[TensorRT] INFO: Tactic -675401754313066228 time 1.12784
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 0) [Convolution] + (Unnamed Layer* 2) [Activation](1)
[TensorRT] INFO: Tactic 0 time 3.46669
[TensorRT] INFO: Tactic 1 scratch requested: 98816, available: 0
[TensorRT] INFO: Tactic 2 scratch requested: 589824, available: 0
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 0) [Convolution] + (Unnamed Layer* 2) [Activation](33)
[TensorRT] INFO: --------------- Chose 3 (0)
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 3) [Convolution] + (Unnamed Layer* 5) [Activation](3)
[TensorRT] INFO: Tactic 0 time 1.8392
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 3) [Convolution] + (Unnamed Layer* 5) [Activation](2)
[TensorRT] INFO: Tactic 6 time 3.87354
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 3) [Convolution] + (Unnamed Layer* 5) [Activation](14)
[TensorRT] INFO: Tactic 3146172331490511787 time 3.48906
[TensorRT] INFO: Tactic 3528302785056538033 time 2.37526
[TensorRT] INFO: Tactic 5443600094180187792 time 2.14822
[TensorRT] INFO: Tactic 5824828673459742858 time 3.40544
[TensorRT] INFO: Tactic -6618588952828687390 time 1.55994
[TensorRT] INFO: Tactic -6362554771847758902 time 1.27136
[TensorRT] INFO: Tactic -2701242286872672544 time 2.89786
[TensorRT] INFO: Tactic -2535759802710599445 time 1.22448
[TensorRT] INFO: Tactic -675401754313066228 time 1.25974
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 3) [Convolution] + (Unnamed Layer* 5) [Activation](1)
[TensorRT] INFO: Tactic 0 time 4.83414
[TensorRT] INFO: Tactic 1 scratch requested: 25088, available: 0
[TensorRT] INFO: Tactic 2 scratch requested: 4718592, available: 0
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 3) [Convolution] + (Unnamed Layer* 5) [Activation](33)
[TensorRT] INFO: --------------- Chose 14 (-2535759802710599445)
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 6) [Convolution] + (Unnamed Layer* 8) [Activation](3)
[TensorRT] INFO: Tactic 0 time 1.11046
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 6) [Convolution] + (Unnamed Layer* 8) [Activation](2)
[TensorRT] INFO: Tactic 6 time 1.98544
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 6) [Convolution] + (Unnamed Layer* 8) [Activation](14)
[TensorRT] INFO: Tactic 3146172331490511787 time 1.13718
[TensorRT] INFO: Tactic 3528302785056538033 time 2.6223
[TensorRT] INFO: Tactic 5443600094180187792 time 1.37126
[TensorRT] INFO: Tactic 5824828673459742858 time 1.11578
[TensorRT] INFO: Tactic -6618588952828687390 time 1.48784
[TensorRT] INFO: Tactic -6362554771847758902 time 1.21818
[TensorRT] INFO: Tactic -2701242286872672544 time 1.12512
[TensorRT] INFO: Tactic -2535759802710599445 time 1.17814
[TensorRT] INFO: Tactic -675401754313066228 time 1.21862
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 6) [Convolution] + (Unnamed Layer* 8) [Activation](1)
[TensorRT] INFO: Tactic 0 time 3.2784
[TensorRT] INFO: Tactic 1 scratch requested: 6656, available: 0
[TensorRT] INFO: Tactic 2 scratch requested: 2359296, available: 0
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 6) [Convolution] + (Unnamed Layer* 8) [Activation](33)
[TensorRT] INFO: --------------- Chose 3 (0)
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 9) [Convolution] + (Unnamed Layer* 11) [Activation](3)
[TensorRT] INFO: Tactic 0 time 1.076
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 9) [Convolution] + (Unnamed Layer* 11) [Activation](2)
[TensorRT] INFO: Tactic 6 time 2.21725
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 9) [Convolution] + (Unnamed Layer* 11) [Activation](14)
[TensorRT] INFO: Tactic 3146172331490511787 time 1.10688
[TensorRT] INFO: Tactic 3528302785056538033 time 1.73837
[TensorRT] INFO: Tactic 5443600094180187792 time 2.2024
[TensorRT] INFO: Tactic 5824828673459742858 time 1.07421
[TensorRT] INFO: Tactic -6618588952828687390 time 2.12707
[TensorRT] INFO: Tactic -6362554771847758902 time 1.1808
[TensorRT] INFO: Tactic -2701242286872672544 time 1.95904
[TensorRT] INFO: Tactic -2535759802710599445 time 1.11232
[TensorRT] INFO: Tactic -675401754313066228 time 1.16397
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 9) [Convolution] + (Unnamed Layer* 11) [Activation](1)
[TensorRT] INFO: Tactic 0 time 2.84109
[TensorRT] INFO: Tactic 1 scratch requested: 2048, available: 0
[TensorRT] INFO: Tactic 2 scratch requested: 1179648, available: 0
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 9) [Convolution] + (Unnamed Layer* 11) [Activation](33)
[TensorRT] INFO: --------------- Chose 14 (5824828673459742858)
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 12) [Convolution] + (Unnamed Layer* 14) [Activation](3)
[TensorRT] INFO: Tactic 0 time 2.10848
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 12) [Convolution] + (Unnamed Layer* 14) [Activation](2)
[TensorRT] INFO: Tactic 6 time 3.44688
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 12) [Convolution] + (Unnamed Layer* 14) [Activation](14)
[TensorRT] INFO: Tactic 3146172331490511787 time 2.18742
[TensorRT] INFO: Tactic 3528302785056538033 time 3.88074
[TensorRT] INFO: Tactic 5443600094180187792 time 4.26378
[TensorRT] INFO: Tactic 5824828673459742858 time 2.11472
[TensorRT] INFO: Tactic -6618588952828687390 time 3.1528
[TensorRT] INFO: Tactic -6362554771847758902 time 2.26848
[TensorRT] INFO: Tactic -2701242286872672544 time 3.07386
[TensorRT] INFO: Tactic -2535759802710599445 time 2.21392
[TensorRT] INFO: Tactic -675401754313066228 time 2.2447
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 12) [Convolution] + (Unnamed Layer* 14) [Activation](1)
[TensorRT] INFO: Tactic 0 time 2.87024
[TensorRT] INFO: Tactic 1 scratch requested: 512, available: 0
[TensorRT] INFO: Tactic 2 scratch requested: 589824, available: 0
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 12) [Convolution] + (Unnamed Layer* 14) [Activation](33)
[TensorRT] INFO: --------------- Chose 3 (0)
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 15) [Deconvolution](4)
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 15) [Deconvolution](16)
[TensorRT] INFO: 
[TensorRT] INFO: --------------- Timing (Unnamed Layer* 15) [Deconvolution](22)
[TensorRT] INFO: Tactic 0 scratch requested: 655360, available: 0
[TensorRT] ERROR: Internal error: could not find any implementation for node (Unnamed Layer* 15) [Deconvolution], try increasing the workspace size with IBuilder::setMaxWorkspaceSize()
[TensorRT] ERROR: ../builder/tacticOptimizer.cpp (1330) - OutOfMemory Error in computeCosts: 0

I know that converter doesn’t implement whole functionality of torch. And i checked all classes used in my model. All except nn.Sequential are supported. As for nn.Sequential - in the example notebook (torch2trt github) for segmentation used deeplabv3_resnet101 model, which contains this layer.

Any suggestions, how can convert my model successfully?

Hi,

The error indicates that you are running out of memory:

[TensorRT] ERROR: ../builder/tacticOptimizer.cpp (1330) - OutOfMemory Error in computeCosts: 0

Is there any parameter can be adjusted to reduce the model size? Like batchsize?

Thanks.

No, batchsize set to minimum value - 2. I’ve even created model with batchsize 1, but got the same error.

Is there way to increase the workspace size for tensorrt?

Found in the source code. Workspace size can be set via parameter during model converting

model_trt = torch2trt(model, [data], max_workspace_size=2<<10)