diff --git a/yolort/models/darknetv4.py b/yolort/models/darknetv4.py index 235cd868..d37cbd7d 100644 --- a/yolort/models/darknetv4.py +++ b/yolort/models/darknetv4.py @@ -4,7 +4,7 @@ import torch from torch import nn, Tensor from yolort.utils import load_state_dict_from_url -from yolort.v5 import Conv, Focus, BottleneckCSP, C3, SPPF +from yolort.v5 import Conv, Focus, BottleneckCSP, C3, SPP from ._utils import _make_divisible @@ -95,7 +95,7 @@ def __init__( # building last CSP blocks last_channel = _make_divisible(last_channel * width_multiple, round_nearest) layers.append(Conv(input_channel, last_channel, k=3, s=2, version=version)) - layers.append(SPPF(last_channel, last_channel, k=5, version=version)) + layers.append(SPP(last_channel, last_channel, k=(5, 9, 13), version=version)) self.features = nn.Sequential(*layers) self.avgpool = nn.AdaptiveAvgPool2d(1) diff --git a/yolort/models/path_aggregation_network.py b/yolort/models/path_aggregation_network.py index b82c4bbb..faa649d4 100644 --- a/yolort/models/path_aggregation_network.py +++ b/yolort/models/path_aggregation_network.py @@ -3,7 +3,7 @@ import torch from torch import nn, Tensor -from yolort.v5 import Conv, BottleneckCSP, C3, SPPF +from yolort.v5 import Conv, BottleneckCSP, C3, SPP class IntermediateLevelP6(nn.Module): @@ -106,7 +106,7 @@ def __init__( depth_gain = max(round(3 * depth_multiple), 1) if version == "r6.0": - init_block = SPPF(in_channels[-1], in_channels[-1], k=5) + init_block = SPP(in_channels[-1], in_channels[-1], k=(5, 9, 13)) elif version in ["r3.1", "r4.0"]: init_block = block(in_channels[-1], in_channels[-1], n=depth_gain, shortcut=False) else: