Replace Focus to Focus2

zhiqwang · Aug 21, 2022 · f949d69 · f949d69
1 parent 07749be
commit f949d69
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/yolort/models/darknetv4.py b/yolort/models/darknetv4.py
@@ -5,7 +5,7 @@
 import torch
 from torch import nn, Tensor
 from yolort.utils import load_state_dict_from_url
-from yolort.v5 import BottleneckCSP, C3, Conv, Focus, SPP
+from yolort.v5 import BottleneckCSP, C3, Conv, Focus2, SPP
 
 from ._utils import _make_divisible
 
@@ -82,7 +82,7 @@ def __init__(
 
         # building first layer
         out_channel = _make_divisible(input_channel * width_multiple, round_nearest)
-        layers.append(Focus(3, out_channel, k=3, version=version))
+        layers.append(Focus2(3, out_channel, k=3, version=version))
         input_channel = out_channel
 
         # building CSP blocks

diff --git a/yolort/v5/models/__init__.py b/yolort/v5/models/__init__.py
@@ -10,6 +10,7 @@
     DWConv,
     Expand,
     Focus,
+    Focus2,
     focus_transform,
     GhostBottleneck,
     GhostConv,
@@ -27,6 +28,7 @@
     "SPPF",
     "DWConv",
     "Focus",
+    "Focus2",
     "BottleneckCSP",
     "C3",
     "C3TR",