yolov5-C2F：将yolov8中C2F结构与yolov5结合 并TensorRT（API）部署

在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<std::string, Weights>& 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