Revert SPPF to SPP in r4.0

zhiqwang · Dec 1, 2021 · 65c1951 · 65c1951
1 parent eea4e4d
commit 65c1951
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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)