new CSP model release
This commit is contained in:
1 parent
854ccb9a74
commit
c14368d768
8 files changed
+54
-53
No files matched your search
+8
-5
@@ -45,14 +45,17 @@ class Detect(nn.Module):
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, model_yaml='yolov5s.yaml', ch=3, nc=None): # model, input channels, number of classes
|
||||
def __init__(self, model_cfg='yolov5s.yaml', ch=3, nc=None): # model, input channels, number of classes
|
||||
super(Model, self).__init__()
|
||||
with open(model_yaml) as f:
|
||||
self.md = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
||||
if nc:
|
||||
self.md['nc'] = nc # override yaml value
|
||||
if type(model_cfg) is dict:
|
||||
self.md = model_cfg # model dict
|
||||
else: # is *.yaml
|
||||
with open(model_cfg) as f:
|
||||
self.md = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
||||
|
||||
# Define model
|
||||
if nc:
|
||||
self.md['nc'] = nc # override yaml value
|
||||
self.model, self.save, ch = parse_model(self.md, ch=[ch]) # model, savelist, ch_out
|
||||
# print([x.shape for x in self.forward(torch.zeros(1, 3, 64, 64))])
|
||||
|
||||
|
||||
+6
-6
@@ -16,29 +16,29 @@ backbone:
|
||||
[-1, 1, Conv, [128, 3, 2]], # 2-P2/4
|
||||
[-1, 3, Bottleneck, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 4-P3/8
|
||||
[-1, 9, Bottleneck, [256]],
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 6-P4/16
|
||||
[-1, 9, Bottleneck, [512]],
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 8-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 3, Bottleneck, [1024]], # 10
|
||||
[-1, 6, BottleneckCSP, [1024]], # 10
|
||||
]
|
||||
|
||||
# yolov5 head
|
||||
head:
|
||||
[[-1, 3, Bottleneck, [1024, False]], # 11
|
||||
[[-1, 3, BottleneckCSP, [1024, False]], # 11
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 12 (P5/32-large)
|
||||
|
||||
[-2, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 3, Bottleneck, [512, False]],
|
||||
[-1, 3, BottleneckCSP, [512, False]],
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 17 (P4/16-medium)
|
||||
|
||||
[-2, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 3, Bottleneck, [256, False]],
|
||||
[-1, 3, BottleneckCSP, [256, False]],
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 22 (P3/8-small)
|
||||
|
||||
[[], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
|
||||
+6
-6
@@ -16,29 +16,29 @@ backbone:
|
||||
[-1, 1, Conv, [128, 3, 2]], # 2-P2/4
|
||||
[-1, 3, Bottleneck, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 4-P3/8
|
||||
[-1, 9, Bottleneck, [256]],
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 6-P4/16
|
||||
[-1, 9, Bottleneck, [512]],
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 8-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 3, Bottleneck, [1024]], # 10
|
||||
[-1, 6, BottleneckCSP, [1024]], # 10
|
||||
]
|
||||
|
||||
# yolov5 head
|
||||
head:
|
||||
[[-1, 3, Bottleneck, [1024, False]], # 11
|
||||
[[-1, 3, BottleneckCSP, [1024, False]], # 11
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 12 (P5/32-large)
|
||||
|
||||
[-2, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 3, Bottleneck, [512, False]],
|
||||
[-1, 3, BottleneckCSP, [512, False]],
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 17 (P4/16-medium)
|
||||
|
||||
[-2, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 3, Bottleneck, [256, False]],
|
||||
[-1, 3, BottleneckCSP, [256, False]],
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 22 (P3/8-small)
|
||||
|
||||
[[], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
|
||||
+6
-6
@@ -16,29 +16,29 @@ backbone:
|
||||
[-1, 1, Conv, [128, 3, 2]], # 2-P2/4
|
||||
[-1, 3, Bottleneck, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 4-P3/8
|
||||
[-1, 9, Bottleneck, [256]],
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 6-P4/16
|
||||
[-1, 9, Bottleneck, [512]],
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 8-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 3, Bottleneck, [1024]], # 10
|
||||
[-1, 6, BottleneckCSP, [1024]], # 10
|
||||
]
|
||||
|
||||
# yolov5 head
|
||||
head:
|
||||
[[-1, 3, Bottleneck, [1024, False]], # 11
|
||||
[[-1, 3, BottleneckCSP, [1024, False]], # 11
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 12 (P5/32-large)
|
||||
|
||||
[-2, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 3, Bottleneck, [512, False]],
|
||||
[-1, 3, BottleneckCSP, [512, False]],
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 17 (P4/16-medium)
|
||||
|
||||
[-2, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 3, Bottleneck, [256, False]],
|
||||
[-1, 3, BottleneckCSP, [256, False]],
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 22 (P3/8-small)
|
||||
|
||||
[[], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
|
||||
+6
-6
@@ -16,29 +16,29 @@ backbone:
|
||||
[-1, 1, Conv, [128, 3, 2]], # 2-P2/4
|
||||
[-1, 3, Bottleneck, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 4-P3/8
|
||||
[-1, 9, Bottleneck, [256]],
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 6-P4/16
|
||||
[-1, 9, Bottleneck, [512]],
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 8-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 3, Bottleneck, [1024]], # 10
|
||||
[-1, 6, BottleneckCSP, [1024]], # 10
|
||||
]
|
||||
|
||||
# yolov5 head
|
||||
head:
|
||||
[[-1, 3, Bottleneck, [1024, False]], # 11
|
||||
[[-1, 3, BottleneckCSP, [1024, False]], # 11
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 12 (P5/32-large)
|
||||
|
||||
[-2, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 3, Bottleneck, [512, False]],
|
||||
[-1, 3, BottleneckCSP, [512, False]],
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 17 (P4/16-medium)
|
||||
|
||||
[-2, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 3, Bottleneck, [256, False]],
|
||||
[-1, 3, BottleneckCSP, [256, False]],
|
||||
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 22 (P3/8-small)
|
||||
|
||||
[[], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
|
||||
Reference in new issue
Block a user