From b6558c17e4cef9545850777fba6ac7b063238339 Mon Sep 17 00:00:00 2001
From: Zhiqiang Wang <zhiqwang@foxmail.com>
Date: Mon, 22 Aug 2022 02:03:51 +0800
Subject: [PATCH] Test Focus2

---
 yolort/models/darknetv4.py   | 4 ++--
 yolort/v5/models/__init__.py | 2 --
 yolort/v5/models/common.py   | 4 ++--
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/yolort/models/darknetv4.py b/yolort/models/darknetv4.py
index 6720d3e5..e9cd00ed 100644
--- 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
index 0b91a1f0..1ee87248 100644
--- 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
index b1a2e708..e4e36893 100644
--- 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