|
|
@@ -11,9 +11,9 @@ class SPPF(nn.Module):
|
|
|
"""
|
|
|
This code referenced to https://github.com/ultralytics/yolov5
|
|
|
"""
|
|
|
- def __init__(self, in_dim, out_dim, expand_ratio=0.5):
|
|
|
+ def __init__(self, in_dim, out_dim, expansion=0.5):
|
|
|
super().__init__()
|
|
|
- inter_dim = int(in_dim * expand_ratio)
|
|
|
+ inter_dim = int(in_dim * expansion)
|
|
|
self.out_dim = out_dim
|
|
|
self.cv1 = ConvModule(in_dim, inter_dim, kernel_size=1)
|
|
|
self.cv2 = ConvModule(inter_dim * 4, out_dim, kernel_size=1)
|
|
|
@@ -30,25 +30,24 @@ class SPPFBlockCSP(nn.Module):
|
|
|
def __init__(self,
|
|
|
in_dim: int,
|
|
|
out_dim: int,
|
|
|
- expand_ratio: float = 0.5,
|
|
|
+ expansion: float = 0.5,
|
|
|
):
|
|
|
super(SPPFBlockCSP, self).__init__()
|
|
|
- inter_dim = int(in_dim * expand_ratio)
|
|
|
+ inter_dim = int(in_dim * expansion)
|
|
|
self.out_dim = out_dim
|
|
|
self.cv1 = ConvModule(in_dim, inter_dim, kernel_size=1)
|
|
|
self.cv2 = ConvModule(in_dim, inter_dim, kernel_size=1)
|
|
|
self.m = nn.Sequential(
|
|
|
ConvModule(inter_dim, inter_dim, kernel_size=3),
|
|
|
- SPPF(inter_dim, inter_dim, expand_ratio=1.0),
|
|
|
+ SPPF(inter_dim, inter_dim, expansion=1.0),
|
|
|
ConvModule(inter_dim, inter_dim, kernel_size=3)
|
|
|
)
|
|
|
self.cv3 = ConvModule(inter_dim * 2, self.out_dim, kernel_size=1)
|
|
|
|
|
|
def forward(self, x):
|
|
|
x1 = self.cv1(x)
|
|
|
- x2 = self.cv2(x)
|
|
|
- x3 = self.m(x2)
|
|
|
- y = self.cv3(torch.cat([x1, x3], dim=1))
|
|
|
+ x2 = self.m(self.cv2(x))
|
|
|
+ y = self.cv3(torch.cat([x1, x2], dim=1))
|
|
|
|
|
|
return y
|
|
|
|
|
|
@@ -59,7 +58,7 @@ if __name__=='__main__':
|
|
|
# Build a neck
|
|
|
in_dim = 512
|
|
|
out_dim = 512
|
|
|
- model = SPPFBlockCSP(512, 512, expand_ratio=0.5)
|
|
|
+ model = SPPFBlockCSP(512, 512, expansion=0.5)
|
|
|
|
|
|
# Randomly generate a input data
|
|
|
x = torch.randn(2, in_dim, 20, 20)
|