Open leaf918 opened 2 years ago
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)
你可以把maxpool3d改成两个maxpool2d