diff --git a/README.md b/README.md
index 6306e55ec866..c6b638003c62 100755
--- a/README.md
+++ b/README.md
@@ -41,9 +41,13 @@ $ pip install -U -r requirements.txt
## Tutorials
* [Notebook](https://github.com/ultralytics/yolov5/blob/master/tutorial.ipynb) 
+* [Kaggle](https://www.kaggle.com/ultralytics/yolov5-tutorial)
* [Train Custom Data](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data)
-* [Google Cloud Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart)
-* [Docker Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart) 
+* [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36)
+* [ONNX and TorchScript Export](https://github.com/ultralytics/yolov5/issues/251)
+* [Test-Time Augmentation (TTA)](https://github.com/ultralytics/yolov5/issues/303)
+* [Google Cloud Quickstart](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart)
+* [Docker Quickstart](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart) 
## Inference
diff --git a/data/get_coco2017.sh b/data/get_coco2017.sh
index fed57473d5c0..03b2c7e89301 100755
--- a/data/get_coco2017.sh
+++ b/data/get_coco2017.sh
@@ -1,7 +1,11 @@
#!/bin/bash
-# Zip coco folder
-# zip -r coco.zip coco
-# tar -czvf coco.tar.gz coco
+# COCO 2017 dataset http://cocodataset.org
+# Download command: bash yolov5/data/get_coco2017.sh
+# Train command: python train.py --data ./data/coco.yaml
+# Dataset should be placed next to yolov5 folder:
+# /parent_folder
+# /coco
+# /yolov5
# Download labels from Google Drive, accepting presented query
filename="coco2017labels.zip"
diff --git a/data/get_voc.sh b/data/get_voc.sh
new file mode 100644
index 000000000000..b7e66d003133
--- /dev/null
+++ b/data/get_voc.sh
@@ -0,0 +1,214 @@
+# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
+# Download command: bash ./data/get_voc.sh
+# Train command: python train.py --data voc.yaml
+# Dataset should be placed next to yolov5 folder:
+# /parent_folder
+# /VOC
+# /yolov5
+
+start=`date +%s`
+
+# handle optional download dir
+if [ -z "$1" ]
+ then
+ # navigate to ~/tmp
+ echo "navigating to ../tmp/ ..."
+ mkdir -p ../tmp
+ cd ../tmp/
+ else
+ # check if is valid directory
+ if [ ! -d $1 ]; then
+ echo $1 "is not a valid directory"
+ exit 0
+ fi
+ echo "navigating to" $1 "..."
+ cd $1
+fi
+
+echo "Downloading VOC2007 trainval ..."
+# Download the data.
+curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar
+echo "Downloading VOC2007 test data ..."
+curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar
+echo "Done downloading."
+
+# Extract data
+echo "Extracting trainval ..."
+tar -xf VOCtrainval_06-Nov-2007.tar
+echo "Extracting test ..."
+tar -xf VOCtest_06-Nov-2007.tar
+echo "removing tars ..."
+rm VOCtrainval_06-Nov-2007.tar
+rm VOCtest_06-Nov-2007.tar
+
+end=`date +%s`
+runtime=$((end-start))
+
+echo "Completed in" $runtime "seconds"
+
+start=`date +%s`
+
+# handle optional download dir
+if [ -z "$1" ]
+ then
+ # navigate to ~/tmp
+ echo "navigating to ../tmp/ ..."
+ mkdir -p ../tmp
+ cd ../tmp/
+ else
+ # check if is valid directory
+ if [ ! -d $1 ]; then
+ echo $1 "is not a valid directory"
+ exit 0
+ fi
+ echo "navigating to" $1 "..."
+ cd $1
+fi
+
+echo "Downloading VOC2012 trainval ..."
+# Download the data.
+curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
+echo "Done downloading."
+
+
+# Extract data
+echo "Extracting trainval ..."
+tar -xf VOCtrainval_11-May-2012.tar
+echo "removing tar ..."
+rm VOCtrainval_11-May-2012.tar
+
+end=`date +%s`
+runtime=$((end-start))
+
+echo "Completed in" $runtime "seconds"
+
+cd ../tmp
+echo "Spliting dataset..."
+python3 - "$@" < train.txt
+cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt > train.all.txt
+
+python3 - "$@" < 1 else n # depth gain
- if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, ConvPlus, BottleneckCSP]:
+ if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]:
c1, c2 = ch[f], args[0]
# Normal
@@ -182,7 +188,7 @@ def parse_model(md, ch): # model_dict, input_channels(3)
# c2 = make_divisible(c2, 8) if c2 != no else c2
args = [c1, c2, *args[1:]]
- if m is BottleneckCSP:
+ if m in [BottleneckCSP, C3]:
args.insert(2, n)
n = 1
elif m is nn.BatchNorm2d:
@@ -198,7 +204,7 @@ def parse_model(md, ch): # model_dict, input_channels(3)
t = str(m)[8:-2].replace('__main__.', '') # module type
np = sum([x.numel() for x in m_.parameters()]) # number params
m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
- print('%3s%15s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print
+ print('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
layers.append(m_)
ch.append(c2)
diff --git a/test.py b/test.py
index 259d44444bcd..1cfae9591287 100644
--- a/test.py
+++ b/test.py
@@ -22,6 +22,7 @@ def test(data,
# Initialize/load model and set device
if model is None:
training = False
+ merge = opt.merge # use Merge NMS
device = torch_utils.select_device(opt.device, batch_size=batch_size)
# Remove previous
@@ -30,11 +31,8 @@ def test(data,
# Load model
google_utils.attempt_download(weights)
- model = torch.load(weights, map_location=device)['model'].float() # load to FP32
- torch_utils.model_info(model)
- model.fuse()
- model.to(device)
- imgsz = check_img_size(imgsz, s=model.model[-1].stride.max()) # check img_size
+ model = torch.load(weights, map_location=device)['model'].float().fuse().to(device) # load to FP32
+ imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
@@ -59,7 +57,6 @@ def test(data,
# Dataloader
if dataloader is None: # not training
- merge = opt.merge # use Merge NMS
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
diff --git a/train.py b/train.py
index eacd265d2bb7..b9b9f083d4d6 100644
--- a/train.py
+++ b/train.py
@@ -85,8 +85,7 @@ def train(hyp, tb_writer, opt, device):
os.remove(f)
# Create model
- model = Model(opt.cfg).to(device)
- assert model.md['nc'] == nc, '%s nc=%g classes but %s nc=%g classes' % (opt.data, nc, opt.cfg, model.md['nc'])
+ model = Model(opt.cfg, nc=data_dict['nc']).to(device)
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
@@ -117,6 +116,9 @@ def train(hyp, tb_writer, opt, device):
optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
+ # Scheduler https://arxiv.org/pdf/1812.01187.pdf
+ lf = lambda x: (((1 + math.cos(x * math.pi / epochs)) / 2) ** 1.0) * 0.9 + 0.1 # cosine
+ scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
print('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
@@ -162,9 +164,7 @@ def train(hyp, tb_writer, opt, device):
if mixed_precision:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0)
- # Scheduler https://arxiv.org/pdf/1812.01187.pdf
- lf = lambda x: (((1 + math.cos(x * math.pi / epochs)) / 2) ** 1.0) * 0.9 + 0.1 # cosine
- scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
+
scheduler.last_epoch = start_epoch - 1 # do not move
# https://discuss.pytorch.org/t/a-problem-occured-when-resuming-an-optimizer/28822
# plot_lr_scheduler(optimizer, scheduler, epochs)
@@ -382,7 +382,6 @@ def train(hyp, tb_writer, opt, device):
if (best_fitness == fi) and not final_epoch:
torch.save(ckpt, best)
del ckpt
-
# end epoch ----------------------------------------------------------------------------------------------------
# end training
diff --git a/utils/torch_utils.py b/utils/torch_utils.py
index 71c8f4c28539..786b01896d50 100644
--- a/utils/torch_utils.py
+++ b/utils/torch_utils.py
@@ -77,16 +77,36 @@ def find_modules(model, mclass=nn.Conv2d):
return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
+def sparsity(model):
+ # Return global model sparsity
+ a, b = 0., 0.
+ for p in model.parameters():
+ a += p.numel()
+ b += (p == 0).sum()
+ return b / a
+
+
+def prune(model, amount=0.3):
+ # Prune model to requested global sparsity
+ import torch.nn.utils.prune as prune
+ print('Pruning model... ', end='')
+ for name, m in model.named_modules():
+ if isinstance(m, nn.Conv2d):
+ prune.l1_unstructured(m, name='weight', amount=amount) # prune
+ prune.remove(m, 'weight') # make permanent
+ print(' %.3g global sparsity' % sparsity(model))
+
+
def fuse_conv_and_bn(conv, bn):
# https://tehnokv.com/posts/fusing-batchnorm-and-conv/
with torch.no_grad():
# init
- fusedconv = torch.nn.Conv2d(conv.in_channels,
- conv.out_channels,
- kernel_size=conv.kernel_size,
- stride=conv.stride,
- padding=conv.padding,
- bias=True)
+ fusedconv = nn.Conv2d(conv.in_channels,
+ conv.out_channels,
+ kernel_size=conv.kernel_size,
+ stride=conv.stride,
+ padding=conv.padding,
+ bias=True).to(conv.weight.device)
# prepare filters
w_conv = conv.weight.clone().view(conv.out_channels, -1)
@@ -94,10 +114,7 @@ def fuse_conv_and_bn(conv, bn):
fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
# prepare spatial bias
- if conv.bias is not None:
- b_conv = conv.bias
- else:
- b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device)
+ b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
@@ -140,8 +157,8 @@ def load_classifier(name='resnet101', n=2):
# Reshape output to n classes
filters = model.fc.weight.shape[1]
- model.fc.bias = torch.nn.Parameter(torch.zeros(n), requires_grad=True)
- model.fc.weight = torch.nn.Parameter(torch.zeros(n, filters), requires_grad=True)
+ model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True)
+ model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True)
model.fc.out_features = n
return model
@@ -176,21 +193,23 @@ class ModelEMA:
"""
def __init__(self, model, decay=0.9999, device=''):
- # make a copy of the model for accumulating moving average of weights
- self.ema = deepcopy(model)
+ # Create EMA
+ self.ema = deepcopy(model.module if is_parallel(model) else model) # FP32 EMA
self.ema.eval()
self.updates = 0 # number of EMA updates
self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs)
self.device = device # perform ema on different device from model if set
if device:
- self.ema.to(device=device)
+ self.ema.to(device)
for p in self.ema.parameters():
p.requires_grad_(False)
def update(self, model):
- self.updates += 1
- d = self.decay(self.updates)
+ # Update EMA parameters
with torch.no_grad():
+ self.updates += 1
+ d = self.decay(self.updates)
+
msd = model.module.state_dict() if hasattr(model, 'module') else model.state_dict()
esd = self.ema.module.state_dict() if hasattr(self.ema, 'module') else self.ema.state_dict()
for k, v in esd.items():
@@ -200,13 +219,13 @@ def update(self, model):
def update_attr(self, model):
# Assign attributes (which may change during training)
- for k in model.__dict__.keys():
+ for k, v in model.__dict__.items():
# TODO: This is uglyy. Custom attributes should have some specific naming strategy.
- if not (k.startswith('_') or k == 'module' or
- isinstance(getattr(model, k), (torch.distributed.ProcessGroupNCCL, torch.distributed.Reducer))):
+ if not (k.startswith('_') or k in ["process_group", "reducer"] or
+ isinstance(v, (torch.distributed.ProcessGroupNCCL, torch.distributed.Reducer))):
try:
- pickle.dumps(getattr(model, k))
+ pickle.dumps(v)
except Exception:
continue
else:
- setattr(self.ema, k, getattr(model, k))
+ setattr(self.ema, k, v)
diff --git a/utils/utils.py b/utils/utils.py
index cef43a1446a2..4673fa5628e1 100755
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -451,7 +451,9 @@ def compute_loss(p, targets, model): # predictions, targets, model
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
# per output
- nt = 0 # targets
+ nt = 0 # number of targets
+ np = len(p) # number of outputs
+ balance = [1.0, 1.0, 1.0]
for i, pi in enumerate(p): # layer index, layer predictions
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
tobj = torch.zeros_like(pi[..., 0]) # target obj
@@ -481,11 +483,12 @@ def compute_loss(p, targets, model): # predictions, targets, model
# with open('targets.txt', 'a') as file:
# [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
- lobj += BCEobj(pi[..., 4], tobj) # obj loss
+ lobj += BCEobj(pi[..., 4], tobj) * balance[i] # obj loss
- lbox *= h['giou']
- lobj *= h['obj']
- lcls *= h['cls']
+ s = 3 / np # output count scaling
+ lbox *= h['giou'] * s
+ lobj *= h['obj'] * s
+ lcls *= h['cls'] * s
bs = tobj.shape[0] # batch size
if red == 'sum':
g = 3.0 # loss gain
@@ -524,16 +527,14 @@ def build_targets(p, targets, model):
a, t = at[j], t.repeat(na, 1, 1)[j] # filter
# overlaps
+ g = 0.5 # offset
gxy = t[:, 2:4] # grid xy
z = torch.zeros_like(gxy)
if style == 'rect2':
- g = 0.2 # offset
j, k = ((gxy % 1. < g) & (gxy > 1.)).T
a, t = torch.cat((a, a[j], a[k]), 0), torch.cat((t, t[j], t[k]), 0)
offsets = torch.cat((z, z[j] + off[0], z[k] + off[1]), 0) * g
-
elif style == 'rect4':
- g = 0.5 # offset
j, k = ((gxy % 1. < g) & (gxy > 1.)).T
l, m = ((gxy % 1. > (1 - g)) & (gxy < (gain[[2, 3]] - 1.))).T
a, t = torch.cat((a, a[j], a[k], a[l], a[m]), 0), torch.cat((t, t[j], t[k], t[l], t[m]), 0)
@@ -780,11 +781,11 @@ def print_results(k):
wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh
# Filter
- i = (wh0 < 4.0).any(1).sum()
+ i = (wh0 < 3.0).any(1).sum()
if i:
print('WARNING: Extremely small objects found. '
- '%g of %g labels are < 4 pixels in width or height.' % (i, len(wh0)))
- wh = wh0[(wh0 >= 4.0).any(1)] # filter > 2 pixels
+ '%g of %g labels are < 3 pixels in width or height.' % (i, len(wh0)))
+ wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
# Kmeans calculation
from scipy.cluster.vq import kmeans