From 3252a68966f831646db1d496edcea1b0922c487a Mon Sep 17 00:00:00 2001 From: Nikita Date: Fri, 4 Oct 2024 13:28:26 +0000 Subject: [PATCH] feat: replace AveragePool layers with Conv2D in DDRNet --- .../nodes/backbones/ddrnet/blocks.py | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/luxonis_train/nodes/backbones/ddrnet/blocks.py b/luxonis_train/nodes/backbones/ddrnet/blocks.py index 59f76b8b..64804b0e 100644 --- a/luxonis_train/nodes/backbones/ddrnet/blocks.py +++ b/luxonis_train/nodes/backbones/ddrnet/blocks.py @@ -24,15 +24,14 @@ def __init__( """A DAPPM branch. @type kernel_size: int - @param kernel_size: The kernel size for the average pooling. - When stride=0, this parameter is omitted, and - AdaptiveAvgPool2d over all the input is performed. + @param kernel_size: The kernel size. When stride=0, this + parameter is omitted, and AdaptiveAvgPool2d over all the + input is performed. @type stride: int - @param stride: Stride for the average pooling. When stride=0, an + @param stride: Stride for the first convolution. When stride=0, AdaptiveAvgPool2d over all the input is performed (output is - 1x1). When stride=1, no average pooling is performed. When - stride>1, average pooling is performed (scaling the input - down and up again). + 1x1). When stride=1, nothing is performed. When stride>1, a + convolution with stride=stride is performed. @type in_channels: int @param in_channels: Number of input channels. @type branch_channels: int @@ -49,8 +48,14 @@ def __init__( down_list.append(nn.AdaptiveAvgPool2d((1, 1))) elif stride > 1: down_list.append( - nn.AvgPool2d( - kernel_size=kernel_size, stride=stride, padding=stride + nn.Conv2d( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=kernel_size, + stride=stride, + padding=stride, + groups=in_channels, + bias=False, ) )