github代码:https://github.com/mary-0830/Yolov5_Widerface
源代码链接:https://github.com/ultralytics/yolov5
训练自己数据集:https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
注意: 若出现fatal: unable to access ‘https://github.comxxxxxxxxxxx’: Failed to connect to xxxxxxxxxxxxx这类错误
解决方法: 将命令行里的http改为git重新执行。 链接
下载链接:http://shuoyang1213.me/WIDERFACE/
使用google下载比较快,也可以用腾讯下载,还要把标注文件给下载下来。
参考链接:https://blog.csdn.net/mary_0830/article/details/116589279
在data/中添加一个widerface.yaml的配置文件,如下图所示,nc改为1,name改为"face",train和val改为之前转好的yolo格式。
将下图参数设置一下,然后执行 python train.py
即可。
开始训练
runtimeError: Couldn't load custom C++ ops. This can happen if your PyTorch and torchvision versions are incompatible, or if you had errors while compiling torchvision from source.
问题,使用下面的方法三解决。找到对应的版本。原链接参考:https://blog.csdn.net/weixin_44936889/article/details/110732476
可以看到原始模型有600MB,而进行模型量化后的模型只有166MB,压缩了很多空间。
源码:
# coding:utf-8
# 模型量化
import os
import torch
import torch
import torch.nn as nn
from tqdm import tqdm
def autopad(k, p=None):
# Pad to 'same'
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
return p
class Conv(nn.Module):
# Standard convolution
# ch_in, ch_out, kernel, stride, padding, groups
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
super(Conv, self).__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p),
groups=g, bias=False)
self.bn = nn.BatchNorm2d(c2)
self.act = nn.Hardswish() if act else nn.Identity()
def forward(self, x):
return self.act(self.bn(self.conv(x)))
def fuseforward(self, x):
return self.act(self.conv(x))
class Ensemble(nn.ModuleList):
# Ensemble of models
def __init__(self):
super(Ensemble, self).__init__()
def forward(self, x, augment=False):
y = []
for module in self:
y.append(module(x, augment)[0])
# y = torch.stack(y).max(0)[0] # max ensemble
# y = torch.cat(y, 1) # nms ensemble
y = torch.stack(y).mean(0) # mean ensemble
return y, None # inference, train output
def attempt_load(weights, map_location=None):
model = Ensemble()
for w in weights if isinstance(weights, list) else [weights]:
# load FP32 model
model.append(torch.load(
w, map_location=map_location)['model'].float().fuse().eval())
# Compatibility updates
for m in tqdm(model.modules()):
if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
m.inplace = True # pytorch 1.7.0 compatibility
elif type(m) is Conv:
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
if len(model) == 1:
return model[-1] # return model
else:
print('Ensemble created with %s\n' % weights)
for k in ['names', 'stride']:
setattr(model, k, getattr(model[-1], k))
return model # return ensemble
def select_device(device='', batch_size=None):
# device = 'cpu' or '0' or '0,1,2,3'
cpu_request = device.lower() == 'cpu'
if device and not cpu_request: # if device requested other than 'cpu'
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
assert torch.cuda.is_available(
), 'CUDA unavailable, invalid device %s requested' % device # check availablity
cuda = False if cpu_request else torch.cuda.is_available()
if cuda:
c = 1024 ** 2 # bytes to MB
ng = torch.cuda.device_count()
if ng > 1 and batch_size: # check that batch_size is compatible with device_count
assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (
batch_size, ng)
x = [torch.cuda.get_device_properties(i) for i in range(ng)]
s = f'Using torch {torch.__version__} '
for i in range(0, ng):
if i == 1:
s = ' ' * len(s)
return torch.device('cuda:0' if cuda else 'cpu')
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--in_weights', type=str,
default='last.pt', help='initial weights path')
parser.add_argument('--out_weights', type=str,
default='slim_model.pt', help='output weights path')
parser.add_argument('--device', type=str, default='0', help='device')
opt = parser.parse_args()
device = select_device(opt.device)
model = attempt_load(opt.in_weights, map_location=device)
model.to(device).eval()
model.half()
torch.save(model, opt.out_weights)
print('done.')
print('-[INFO] before: {} kb, after: {} kb'.format(
os.path.getsize(opt.in_weights), os.path.getsize(opt.out_weights)))
在命令行执行:
python slim.py --in_weights best.pt --out_weights slim_model.pt --device 0
笔者的话: 不得不说,Yolov5这个代码维护越来越棒了,方便小白训练。一直在用20年8月的版本,那个版本还是比较简单的,也没有两张显卡训练,也没有其他的日志保存,缺少的信息还是很多的。快一年过去了,这个版本已经非常成熟了,也用到了各个检测任务上,非常棒!~