maxpool3d 在tensorrt中不支持

[Tomhardy13]

你可以把maxpool3d改成两个maxpool2d

[yanik-porto]

For converting to tensorrt, you can change MaxPool3d with this layer :

class MaxPool2dStep(nn.Module):

    def __init__(self, kernel_size, stride, step, stop):
        super(MaxPool2dStep, self).__init__()
        self.step = step
        self.stop = stop
        self.block = nn.Sequential(
            nn.MaxPool2d(kernel_size, stride=stride)
        )
    def forward(self, x):
        device = x.get_device()
        return self.block(x.index_select(1, torch.tensor(range(0, self.stop, self.step)).to(device)))

For example : nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(2, 1, 2)) becomes : MaxPool2dStep(kernel_size=(3, 3), stride=(1, 2), step=2, stop=128)