From a0bb5b4a44a9179d071939267f4f7506408fb324 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Fri, 14 Jun 2024 13:39:31 -0700 Subject: [PATCH] Missing stem_kernel_size argument in EfficientNetFeatures --- timm/models/efficientnet.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/timm/models/efficientnet.py b/timm/models/efficientnet.py index 9f6b6b54..00903d1e 100644 --- a/timm/models/efficientnet.py +++ b/timm/models/efficientnet.py @@ -79,8 +79,8 @@ class EfficientNet(nn.Module): num_features: int = 1280, in_chans: int = 3, stem_size: int = 32, - fix_stem: bool = False, stem_kernel_size: int = 3, + fix_stem: bool = False, output_stride: int = 32, pad_type: str = '', act_layer: Optional[LayerType] = None, @@ -278,6 +278,7 @@ class EfficientNetFeatures(nn.Module): feature_location: str = 'bottleneck', in_chans: int = 3, stem_size: int = 32, + stem_kernel_size: int = 3, fix_stem: bool = False, output_stride: int = 32, pad_type: str = '', @@ -300,7 +301,7 @@ class EfficientNetFeatures(nn.Module): # Stem if not fix_stem: stem_size = round_chs_fn(stem_size) - self.conv_stem = create_conv2d(in_chans, stem_size, 3, stride=2, padding=pad_type) + self.conv_stem = create_conv2d(in_chans, stem_size, stem_kernel_size, stride=2, padding=pad_type) self.bn1 = norm_act_layer(stem_size, inplace=True) # Middle stages (IR/ER/DS Blocks)