Test Focus2

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

Hide file tree

Showing 3 changed files with 4 additions and 6 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, Focus2, SPP
+from yolort.v5 import BottleneckCSP, C3, Conv, Focus, 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(Focus2(3, out_channel, k=3, version=version))
+        layers.append(Focus(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,7 +10,6 @@
     DWConv,
     Expand,
     Focus,
-    Focus2,
     focus_transform,
     GhostBottleneck,
     GhostConv,
@@ -28,7 +27,6 @@
     "SPPF",
     "DWConv",
     "Focus",
-    "Focus2",
     "BottleneckCSP",
     "C3",
     "C3TR",

diff --git a/yolort/v5/models/common.py b/yolort/v5/models/common.py
@@ -207,7 +207,7 @@ def forward(self, x):
         return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
 
 
-class Focus(nn.Module):
+class FocusTransform(nn.Module):
     """
     Focus wh information into c-space
 
@@ -249,7 +249,7 @@ def space_to_depth(x: Tensor) -> Tensor:
     return y
 
 
-class Focus2(nn.Module):
+class Focus(nn.Module):
     """
     Is the Focus layer equivalent to a simple Conv layer?
     https://github.com/ultralytics/yolov5/issues/4825#issue-998038464