search

Tianxiaomo / pytorch-YOLOv4

PyTorch ,ONNX and TensorRT implementation of YOLOv4
Apache License 2.0
4.47k stars 1.49k forks source link

Max pooling issue #12

Closed ersheng-ai closed 4 years ago

ersheng-ai commented 4 years ago

Sorry for bothering you again :-) nn.MaxPool2d is already enough to handle your cases. In line 237 of darknet2pytorch.py Change

            elif block['type'] == 'maxpool':
                pool_size = int(block['size'])
                stride = int(block['stride'])
                if stride > 1:
                    model = nn.MaxPool2d(pool_size, stride)
                else:
                    model = MaxPoolStride1(pool_size)

to

            elif block['type'] == 'maxpool':
                pool_size = int(block['size'])
                stride = int(block['stride'])
                model = nn.MaxPool2d(kernel_size=pool_size, stride=stride, padding=pool_size//2)
Tianxiaomo commented 4 years ago

Thanks for your suggestion, it has been modified.