在yolov5中添加最新yolov8的C2F模块,使用VisDrone2019数据训练测试。在wang-xinyu基础上添加C2F模块(TensorRT API实现)并进行测试;
1、yolov5-C2F模型训练
C3全部改为C2f,其他保持不变。
# YOLOv5 by Ultralytics, GPL-3.0 license
# Parameters
#nc: 80 # number of classes
nc: 10 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
# [-1, 3, C3, [128]],
[-1, 3, C2f, [128, True]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
# [-1, 6, C3, [256]],
[-1, 6, C2f, [256, True]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
# [-1, 9, C3, [512]],
[-1, 6, C2f, [512, True]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
# [-1, 3, C3, [1024]],
[-1, 3, C2f, [1024, True]],
[-1, 1, SPPF, [1024, 5]], # 9
]
# YOLOv5 v6.0 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
# [-1, 3, C3, [512, False]], # 13
[-1, 3, C2f, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
# [-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 3, C2f, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
# [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 3, C2f, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
# [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[-1, 3, C2f, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
在common.py中添加c2f类;
class C2f(nn.Module):
# CSP Bottleneck with 2 convolutions
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__()
self.c = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, 2 * self.c, 1, 1)
self.cv2 = Conv((2 + n) * self.c, c2, 1) # optional act=FReLU(c2)
self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
def forward(self, x):
y = list(self.cv1(x).split((self.c, self.c), 1))
y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1))
yolov5s-c2f.yaml模型修训练60epoch后的模型,其精度一般,训练数据使用VisDrone2019,训练结果如下图:
2、yolov5-C2F模型 Tensorrt部署
按照wang-xinyu的步骤转换即可。
c2f 部署核心代码如下:具体参考:tensorrt-nyy: yolov TensorRT部署;yolov5-6.0网络添加小目标检测头训练部署 ;yolov5-7.0网络添加C2F模块训练部署 - Gitee.com
ILayer* C2f(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname)
{
int c_ = (int)((float)c2 * e);
auto cv1 = convBlock(network, weightMap, input, 2*c_, 1, 1, 1, lname + ".cv1");
Dims d = cv1->getOutput(0)->getDimensions();
ISliceLayer *s1 = network->addSlice(*cv1->getOutput(0), Dims3{ 0, 0, 0 }, Dims3{ d.d[0] / 2, d.d[1], d.d[2] }, Dims3{ 1, 1 ,1});
ISliceLayer *s2 = network->addSlice(*cv1->getOutput(0), Dims3{ d.d[0] / 2, 0, 0 }, Dims3{ d.d[0] / 2, d.d[1], d.d[2] }, Dims3{ 1, 1 ,1});
ITensor* inputTensors_2[] = {s1->getOutput(0), s2->getOutput(0)};
auto cat= network->addConcatenation(inputTensors_2, 2);
ITensor *y1 = s2->getOutput(0);
for (int i = 0; i < n; i++) {
auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i));
y1 = b->getOutput(0);
ITensor* inputTensors[] = { cat->getOutput(0) , b->getOutput(0) };
cat= network->addConcatenation(inputTensors, 2);
}
auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2");
std::cout<<"c2f end"<<"\n";
return cv2;
}
常用的命令
python export.py --weights yolov5s-c2f60.pt --img 640 --batch 1 --dynamic --include=onnx
python gen_wts.py -w=yolov5s-c2f60.pt -o=yolov5s-c2f60.wts
./yolov5_det -s yolov5s-c2f60.wts yolov5s-c2f60.engine s
./yolov5_det -d yolov5s-c2f60.engine ../samples