Commit
•
a814720
1
Parent(s):
1e25775
PyTorch Hub updates
Browse files- README.md +1 -1
- hubconf.py +90 -0
- models/common.py +0 -37
- models/experimental.py +38 -1
- models/yolo.py +5 -5
- models/yolov3-spp_csp.yaml +0 -55
README.md
CHANGED
@@ -87,7 +87,7 @@ To access an up-to-date working environment (with all dependencies including CUD
|
|
87 |
|
88 |
- **GCP** Deep Learning VM with $300 free credit offer: See our [GCP Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart)
|
89 |
- **Google Colab Notebook** with 12 hours of free GPU time. <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
|
90 |
-
- **Docker Image** https://hub.docker.com/r/ultralytics/yolov5. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart)
|
91 |
|
92 |
|
93 |
## Citation
|
|
|
87 |
|
88 |
- **GCP** Deep Learning VM with $300 free credit offer: See our [GCP Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart)
|
89 |
- **Google Colab Notebook** with 12 hours of free GPU time. <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
|
90 |
+
- **Docker Image** https://hub.docker.com/r/ultralytics/yolov5. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart) ![Docker Pulls](https://img.shields.io/docker/pulls/ultralytics/yolov5?logo=docker)
|
91 |
|
92 |
|
93 |
## Citation
|
hubconf.py
ADDED
@@ -0,0 +1,90 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""File for accessing YOLOv5 via PyTorch Hub https://pytorch.org/hub/
|
2 |
+
|
3 |
+
Usage:
|
4 |
+
import torch
|
5 |
+
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, channels=3, classes=80)
|
6 |
+
"""
|
7 |
+
|
8 |
+
dependencies = ['torch', 'pyyaml']
|
9 |
+
import torch
|
10 |
+
|
11 |
+
from models.yolo import Model
|
12 |
+
from utils import google_utils
|
13 |
+
|
14 |
+
|
15 |
+
def create(name, pretrained, channels, classes):
|
16 |
+
"""Creates a specified YOLOv5 model
|
17 |
+
|
18 |
+
Arguments:
|
19 |
+
name (str): name of model, i.e. 'yolov5s'
|
20 |
+
pretrained (bool): load pretrained weights into the model
|
21 |
+
channels (int): number of input channels
|
22 |
+
classes (int): number of model classes
|
23 |
+
|
24 |
+
Returns:
|
25 |
+
pytorch model
|
26 |
+
"""
|
27 |
+
model = Model('models/%s.yaml' % name, channels, classes)
|
28 |
+
if pretrained:
|
29 |
+
ckpt = '%s.pt' % name # checkpoint filename
|
30 |
+
google_utils.attempt_download(ckpt) # download if not found locally
|
31 |
+
state_dict = torch.load(ckpt)['model'].state_dict()
|
32 |
+
state_dict = {k: v for k, v in state_dict if model.state_dict()[k].numel() == v.numel()} # filter
|
33 |
+
model.load_state_dict(state_dict, strict=False) # load
|
34 |
+
return model
|
35 |
+
|
36 |
+
|
37 |
+
def yolov5s(pretrained=False, channels=3, classes=80):
|
38 |
+
"""YOLOv5-small model from https://github.com/ultralytics/yolov5
|
39 |
+
|
40 |
+
Arguments:
|
41 |
+
pretrained (bool): load pretrained weights into the model, default=False
|
42 |
+
channels (int): number of input channels, default=3
|
43 |
+
classes (int): number of model classes, default=80
|
44 |
+
|
45 |
+
Returns:
|
46 |
+
pytorch model
|
47 |
+
"""
|
48 |
+
return create('yolov5s', pretrained, channels, classes)
|
49 |
+
|
50 |
+
|
51 |
+
def yolov5m(pretrained=False, channels=3, classes=80):
|
52 |
+
"""YOLOv5-medium model from https://github.com/ultralytics/yolov5
|
53 |
+
|
54 |
+
Arguments:
|
55 |
+
pretrained (bool): load pretrained weights into the model, default=False
|
56 |
+
channels (int): number of input channels, default=3
|
57 |
+
classes (int): number of model classes, default=80
|
58 |
+
|
59 |
+
Returns:
|
60 |
+
pytorch model
|
61 |
+
"""
|
62 |
+
return create('yolov5m', pretrained, channels, classes)
|
63 |
+
|
64 |
+
|
65 |
+
def yolov5l(pretrained=False, channels=3, classes=80):
|
66 |
+
"""YOLOv5-large model from https://github.com/ultralytics/yolov5
|
67 |
+
|
68 |
+
Arguments:
|
69 |
+
pretrained (bool): load pretrained weights into the model, default=False
|
70 |
+
channels (int): number of input channels, default=3
|
71 |
+
classes (int): number of model classes, default=80
|
72 |
+
|
73 |
+
Returns:
|
74 |
+
pytorch model
|
75 |
+
"""
|
76 |
+
return create('yolov5l', pretrained, channels, classes)
|
77 |
+
|
78 |
+
|
79 |
+
def yolov5x(pretrained=False, channels=3, classes=80):
|
80 |
+
"""YOLOv5-xlarge model from https://github.com/ultralytics/yolov5
|
81 |
+
|
82 |
+
Arguments:
|
83 |
+
pretrained (bool): load pretrained weights into the model, default=False
|
84 |
+
channels (int): number of input channels, default=3
|
85 |
+
classes (int): number of model classes, default=80
|
86 |
+
|
87 |
+
Returns:
|
88 |
+
pytorch model
|
89 |
+
"""
|
90 |
+
return create('yolov5x', pretrained, channels, classes)
|
models/common.py
CHANGED
@@ -1,8 +1,6 @@
|
|
1 |
# This file contains modules common to various models
|
2 |
|
3 |
|
4 |
-
import torch.nn.functional as F
|
5 |
-
|
6 |
from utils.utils import *
|
7 |
|
8 |
|
@@ -58,17 +56,6 @@ class BottleneckCSP(nn.Module):
|
|
58 |
return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
|
59 |
|
60 |
|
61 |
-
class ConvPlus(nn.Module):
|
62 |
-
# Plus-shaped convolution
|
63 |
-
def __init__(self, c1, c2, k=3, s=1, g=1, bias=True): # ch_in, ch_out, kernel, stride, groups
|
64 |
-
super(ConvPlus, self).__init__()
|
65 |
-
self.cv1 = nn.Conv2d(c1, c2, (k, 1), s, (k // 2, 0), groups=g, bias=bias)
|
66 |
-
self.cv2 = nn.Conv2d(c1, c2, (1, k), s, (0, k // 2), groups=g, bias=bias)
|
67 |
-
|
68 |
-
def forward(self, x):
|
69 |
-
return self.cv1(x) + self.cv2(x)
|
70 |
-
|
71 |
-
|
72 |
class SPP(nn.Module):
|
73 |
# Spatial pyramid pooling layer used in YOLOv3-SPP
|
74 |
def __init__(self, c1, c2, k=(5, 9, 13)):
|
@@ -107,27 +94,3 @@ class Concat(nn.Module):
|
|
107 |
|
108 |
def forward(self, x):
|
109 |
return torch.cat(x, self.d)
|
110 |
-
|
111 |
-
|
112 |
-
class MixConv2d(nn.Module):
|
113 |
-
# Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
|
114 |
-
def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
|
115 |
-
super(MixConv2d, self).__init__()
|
116 |
-
groups = len(k)
|
117 |
-
if equal_ch: # equal c_ per group
|
118 |
-
i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices
|
119 |
-
c_ = [(i == g).sum() for g in range(groups)] # intermediate channels
|
120 |
-
else: # equal weight.numel() per group
|
121 |
-
b = [c2] + [0] * groups
|
122 |
-
a = np.eye(groups + 1, groups, k=-1)
|
123 |
-
a -= np.roll(a, 1, axis=1)
|
124 |
-
a *= np.array(k) ** 2
|
125 |
-
a[0] = 1
|
126 |
-
c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
|
127 |
-
|
128 |
-
self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
|
129 |
-
self.bn = nn.BatchNorm2d(c2)
|
130 |
-
self.act = nn.LeakyReLU(0.1, inplace=True)
|
131 |
-
|
132 |
-
def forward(self, x):
|
133 |
-
return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
|
|
|
1 |
# This file contains modules common to various models
|
2 |
|
3 |
|
|
|
|
|
4 |
from utils.utils import *
|
5 |
|
6 |
|
|
|
56 |
return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
|
57 |
|
58 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
59 |
class SPP(nn.Module):
|
60 |
# Spatial pyramid pooling layer used in YOLOv3-SPP
|
61 |
def __init__(self, c1, c2, k=(5, 9, 13)):
|
|
|
94 |
|
95 |
def forward(self, x):
|
96 |
return torch.cat(x, self.d)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
models/experimental.py
CHANGED
@@ -2,7 +2,7 @@ from models.common import *
|
|
2 |
|
3 |
|
4 |
class Sum(nn.Module):
|
5 |
-
#
|
6 |
def __init__(self, n, weight=False): # n: number of inputs
|
7 |
super(Sum, self).__init__()
|
8 |
self.weight = weight # apply weights boolean
|
@@ -23,6 +23,7 @@ class Sum(nn.Module):
|
|
23 |
|
24 |
|
25 |
class GhostConv(nn.Module):
|
|
|
26 |
def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
|
27 |
super(GhostConv, self).__init__()
|
28 |
c_ = c2 // 2 # hidden channels
|
@@ -35,6 +36,7 @@ class GhostConv(nn.Module):
|
|
35 |
|
36 |
|
37 |
class GhostBottleneck(nn.Module):
|
|
|
38 |
def __init__(self, c1, c2, k, s):
|
39 |
super(GhostBottleneck, self).__init__()
|
40 |
c_ = c2 // 2
|
@@ -46,3 +48,38 @@ class GhostBottleneck(nn.Module):
|
|
46 |
|
47 |
def forward(self, x):
|
48 |
return self.conv(x) + self.shortcut(x)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
2 |
|
3 |
|
4 |
class Sum(nn.Module):
|
5 |
+
# Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
|
6 |
def __init__(self, n, weight=False): # n: number of inputs
|
7 |
super(Sum, self).__init__()
|
8 |
self.weight = weight # apply weights boolean
|
|
|
23 |
|
24 |
|
25 |
class GhostConv(nn.Module):
|
26 |
+
# Ghost Convolution https://github.com/huawei-noah/ghostnet
|
27 |
def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
|
28 |
super(GhostConv, self).__init__()
|
29 |
c_ = c2 // 2 # hidden channels
|
|
|
36 |
|
37 |
|
38 |
class GhostBottleneck(nn.Module):
|
39 |
+
# Ghost Bottleneck https://github.com/huawei-noah/ghostnet
|
40 |
def __init__(self, c1, c2, k, s):
|
41 |
super(GhostBottleneck, self).__init__()
|
42 |
c_ = c2 // 2
|
|
|
48 |
|
49 |
def forward(self, x):
|
50 |
return self.conv(x) + self.shortcut(x)
|
51 |
+
|
52 |
+
|
53 |
+
class ConvPlus(nn.Module):
|
54 |
+
# Plus-shaped convolution
|
55 |
+
def __init__(self, c1, c2, k=3, s=1, g=1, bias=True): # ch_in, ch_out, kernel, stride, groups
|
56 |
+
super(ConvPlus, self).__init__()
|
57 |
+
self.cv1 = nn.Conv2d(c1, c2, (k, 1), s, (k // 2, 0), groups=g, bias=bias)
|
58 |
+
self.cv2 = nn.Conv2d(c1, c2, (1, k), s, (0, k // 2), groups=g, bias=bias)
|
59 |
+
|
60 |
+
def forward(self, x):
|
61 |
+
return self.cv1(x) + self.cv2(x)
|
62 |
+
|
63 |
+
|
64 |
+
class MixConv2d(nn.Module):
|
65 |
+
# Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
|
66 |
+
def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
|
67 |
+
super(MixConv2d, self).__init__()
|
68 |
+
groups = len(k)
|
69 |
+
if equal_ch: # equal c_ per group
|
70 |
+
i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices
|
71 |
+
c_ = [(i == g).sum() for g in range(groups)] # intermediate channels
|
72 |
+
else: # equal weight.numel() per group
|
73 |
+
b = [c2] + [0] * groups
|
74 |
+
a = np.eye(groups + 1, groups, k=-1)
|
75 |
+
a -= np.roll(a, 1, axis=1)
|
76 |
+
a *= np.array(k) ** 2
|
77 |
+
a[0] = 1
|
78 |
+
c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
|
79 |
+
|
80 |
+
self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
|
81 |
+
self.bn = nn.BatchNorm2d(c2)
|
82 |
+
self.act = nn.LeakyReLU(0.1, inplace=True)
|
83 |
+
|
84 |
+
def forward(self, x):
|
85 |
+
return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
|
models/yolo.py
CHANGED
@@ -2,7 +2,7 @@ import argparse
|
|
2 |
|
3 |
import yaml
|
4 |
|
5 |
-
from models.
|
6 |
|
7 |
|
8 |
class Detect(nn.Module):
|
@@ -56,12 +56,12 @@ class Model(nn.Module):
|
|
56 |
# Define model
|
57 |
if nc:
|
58 |
self.md['nc'] = nc # override yaml value
|
59 |
-
self.model, self.save
|
60 |
-
# print([x.shape for x in self.forward(torch.zeros(1,
|
61 |
|
62 |
# Build strides, anchors
|
63 |
m = self.model[-1] # Detect()
|
64 |
-
m.stride = torch.tensor([64 / x.shape[-2] for x in self.forward(torch.zeros(1,
|
65 |
m.anchors /= m.stride.view(-1, 1, 1)
|
66 |
self.stride = m.stride
|
67 |
|
@@ -200,7 +200,7 @@ def parse_model(md, ch): # model_dict, input_channels(3)
|
|
200 |
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
|
201 |
layers.append(m_)
|
202 |
ch.append(c2)
|
203 |
-
return nn.Sequential(*layers), sorted(save)
|
204 |
|
205 |
|
206 |
if __name__ == '__main__':
|
|
|
2 |
|
3 |
import yaml
|
4 |
|
5 |
+
from models.experimental import *
|
6 |
|
7 |
|
8 |
class Detect(nn.Module):
|
|
|
56 |
# Define model
|
57 |
if nc:
|
58 |
self.md['nc'] = nc # override yaml value
|
59 |
+
self.model, self.save = parse_model(self.md, ch=[ch]) # model, savelist, ch_out
|
60 |
+
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
61 |
|
62 |
# Build strides, anchors
|
63 |
m = self.model[-1] # Detect()
|
64 |
+
m.stride = torch.tensor([64 / x.shape[-2] for x in self.forward(torch.zeros(1, ch, 64, 64))]) # forward
|
65 |
m.anchors /= m.stride.view(-1, 1, 1)
|
66 |
self.stride = m.stride
|
67 |
|
|
|
200 |
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
|
201 |
layers.append(m_)
|
202 |
ch.append(c2)
|
203 |
+
return nn.Sequential(*layers), sorted(save)
|
204 |
|
205 |
|
206 |
if __name__ == '__main__':
|
models/yolov3-spp_csp.yaml
DELETED
@@ -1,55 +0,0 @@
|
|
1 |
-
# parameters
|
2 |
-
nc: 80 # number of classes
|
3 |
-
depth_multiple: 1.0 # expand model depth
|
4 |
-
width_multiple: 1.0 # expand layer channels
|
5 |
-
|
6 |
-
# anchors
|
7 |
-
anchors:
|
8 |
-
- [10,13, 16,30, 33,23] # P3/8
|
9 |
-
- [30,61, 62,45, 59,119] # P4/16
|
10 |
-
- [116,90, 156,198, 373,326] # P5/32
|
11 |
-
|
12 |
-
# darknet53 backbone
|
13 |
-
backbone:
|
14 |
-
# [from, number, module, args]
|
15 |
-
[[-1, 1, Conv, [32, 3, 1]], # 0
|
16 |
-
[-1, 1, Conv, [64, 3, 2]], # 1-P1/2
|
17 |
-
[-1, 1, BottleneckCSP, [64]],
|
18 |
-
[-1, 1, Conv, [128, 3, 2]], # 3-P2/4
|
19 |
-
[-1, 2, BottleneckCSP, [128]],
|
20 |
-
[-1, 1, Conv, [256, 3, 2]], # 5-P3/8
|
21 |
-
[-1, 8, BottleneckCSP, [256]],
|
22 |
-
[-1, 1, Conv, [512, 3, 2]], # 7-P4/16
|
23 |
-
[-1, 8, BottleneckCSP, [512]],
|
24 |
-
[-1, 1, Conv, [1024, 3, 2]], # 9-P5/32
|
25 |
-
[-1, 4, BottleneckCSP, [1024]], # 10
|
26 |
-
]
|
27 |
-
|
28 |
-
# yolov3-spp head
|
29 |
-
# na = len(anchors[0])
|
30 |
-
head:
|
31 |
-
[[-1, 1, Bottleneck, [1024, False]], # 11
|
32 |
-
[-1, 1, SPP, [512, [5, 9, 13]]],
|
33 |
-
[-1, 1, Conv, [1024, 3, 1]],
|
34 |
-
[-1, 1, Conv, [512, 1, 1]],
|
35 |
-
[-1, 1, Conv, [1024, 3, 1]],
|
36 |
-
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1]], # 16 (P5/32-large)
|
37 |
-
|
38 |
-
[-3, 1, Conv, [256, 1, 1]],
|
39 |
-
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
40 |
-
[[-1, 8], 1, Concat, [1]], # cat backbone P4
|
41 |
-
[-1, 1, Bottleneck, [512, False]],
|
42 |
-
[-1, 1, Bottleneck, [512, False]],
|
43 |
-
[-1, 1, Conv, [256, 1, 1]],
|
44 |
-
[-1, 1, Conv, [512, 3, 1]],
|
45 |
-
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1]], # 24 (P4/16-medium)
|
46 |
-
|
47 |
-
[-3, 1, Conv, [128, 1, 1]],
|
48 |
-
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
49 |
-
[[-1, 6], 1, Concat, [1]], # cat backbone P3
|
50 |
-
[-1, 1, Bottleneck, [256, False]],
|
51 |
-
[-1, 2, Bottleneck, [256, False]],
|
52 |
-
[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1]], # 30 (P3/8-small)
|
53 |
-
|
54 |
-
[[], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
55 |
-
]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|