libtorch (pytorch c++) 教程（七）

最新推荐文章于 2024-05-23 09:43:14 发布

AllentDan

最新推荐文章于 2024-05-23 09:43:14 发布

阅读量1.1k

点赞数 5

分类专栏： libtorch 文章标签：深度学习人工智能计算机视觉 pytorch c++

本文链接：https://blog.csdn.net/allentdan/article/details/116429281

版权

本教程介绍了如何使用C++和libtorch实现YOLOv4-tiny目标检测模型，包括模型结构、骨干网络CSPDarknet53-tiny的搭建，以及预测和训练的流程。模型以CSPDarknet53-tiny为骨干网络，FPN为颈部，Yolo head为头部，提供训练和预测功能，具有速度优势。

摘要由CSDN通过智能技术生成

阅读本文需要有基础的pytorch编程经验，目标检测框架相关知识，不用很深入，大致了解概念即可。

本章简要介绍如何如何用C++实现一个目标检测器模型，该模型具有训练和预测的功能。本文的分割模型架构使用yolov4-tiny结构，代码结构参考了bubbliiiing yolov4-tiny，本文分享的c++模型几乎完美复现了pytorch的版本，且具有速度优势，30-40%的速度提升。

模型简介

简单介绍一下yolov4-tiny模型。yolov4-tiny模型是YOLO(you only look once)系列模型中，version 4的轻巧版，相比于yolov4，它牺牲了部分精度以实现速度上的大幅提升。yolov4_tiny模型结构如图（图片来源自这）：
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-XGhpI22W-1620215360286)(https://raw.githubusercontent.com/AllentDan/ImageBase/main/detection/yolo4_tiny.png)]
可以发现模型结构非常简单，以CSPDarknet53-tiny为骨干网络，FPN为颈部(neck)，Yolo head为头部。最后输出两个特征层，分别是原图下采样32倍和下采样16倍的特征图。训练时，以这两个特征图分别输入损失计算中计算损失，再将损失求和(或平均，怎么都好)，后做反向传播，预测时将两个特征图解码出的结果做并集再做NMS(非极大值抑制)。

骨干网络

CSPDarknet53-tiny是CSPNet的一种，CSPNet发表于CVPR2019，是用于提升目标检测模型检测性能的一种骨干网络。感兴趣的同学可以去看原文，简单理解该论文贡献，就是将特征层沿着通道维度切成两片，两片分别做不同的卷积，然后再拼接起来，这样做相比于直接对原图做特征提取，能减少计算量。

默认看过我的libtorch系列教程的前部分,直接上代码。首先是基本单元，由Conv2d + BatchNorm2d + LeakyReLU构成。

//Conv2d + BatchNorm2d + LeakyReLU
class BasicConvImpl : public torch::nn::Module {
   
public:
	BasicConvImpl(int in_channels, int out_channels, int kernel_size, int stride = 1);
	torch::Tensor forward(torch::Tensor x);
private:
	// Declare layers
	torch::nn::Conv2d conv{
    nullptr };
	torch::nn::BatchNorm2d bn{
    nullptr };
	torch::nn::LeakyReLU acitivation{
    nullptr };
}; TORCH_MODULE(BasicConv);

BasicConvImpl::BasicConvImpl(int in_channels, int out_channels, int kernel_size, 
	int stride) :
	conv(conv_options(in_channels, out_channels, kernel_size, stride, 
		int(kernel_size / 2), 1, false)),
	bn(torch::nn::BatchNorm2d(out_channels)),
	acitivation(torch::nn::LeakyReLU(torch::nn::LeakyReLUOptions().negative_slope(0.1)))
{
   
	register_module("conv", conv);
	register_module("bn", bn);
}

torch::Tensor BasicConvImpl::forward(torch::Tensor x)
{
   
	x = conv->forward(x);
	x = bn->forward(x);
	x = acitivation(x);
	return x;
}

该层作为基本模块，将在后期作为搭积木的基本块，搭建yolo4_tiny。

然后是Resblock_body模块，

class Resblock_bodyImpl : public torch::nn::Module {
   
public:
	Resblock_bodyImpl(int in_channels, int out_channels);
	std::vector<torch::Tensor> forward(torch::Tensor x);
private:
	int out_channels;
	BasicConv conv1{
    nullptr };
	BasicConv conv2{
    nullptr };
	BasicConv conv3{
    nullptr };
	BasicConv conv4{
    nullptr };
	torch::nn::MaxPool2d maxpool{
    nullptr };
}; TORCH_MODULE(Resblock_body);

Resblock_bodyImpl::Resblock_bodyImpl(int in_channels, int out_channels) {
   
	this->out_channels = out_channels;
	conv1 = BasicConv(in_channels, out_channels, 3);
	conv2 = BasicConv(out_channels / 2, out_channels / 2, 3);
	conv3 = BasicConv(out_channels / 2, out_channels / 2, 3);
	conv4 = BasicConv(out_channels, out_channels, 1);
	maxpool = torch::nn::MaxPool2d(maxpool_options(2, 2));

	register_module("conv1", conv1);
	register_module("conv2", conv2);
	register_module("conv3", conv3);
	register_module("conv4", conv4);

}
std::vector<torch::Tensor> Resblock_bodyImpl::forward(torch::Tensor x) {
   
	auto c = out_channels;
	x = conv1->forward(x);
	auto route = x;

	x = torch::split(x, c / 2, 1)[1];
	x = conv2->forward(x);
	auto route1 = x;

	x = conv3->forward(x);
	x = torch::cat({
    x, route1 }, 1);
	x = conv4->forward(x);
	auto feat = x;

	x = torch::cat({
    route, x }, 1);
	x = maxpool->forward(x);
	return std::vector<torch::Tensor>({
    x,feat });
}

最后是骨干网络主体

class CSPdarknet53_tinyImpl : public torch::nn::Module
{
   
public:
	CSPdarknet53_tinyImpl();
	std::vector<torch::Tensor> forward(torch::Tensor x);
private:
	BasicConv conv1{
    nullptr };
	BasicConv conv2{
    nullptr };
	Resblock_body resblock_body1{
    nullptr };
	Resblock_body resblock_body2{
    nullptr };
	Resblock_body resblock_body3{
    nullptr };
	BasicConv conv3{
    nullptr };
	int num_features = 1;
}; TORCH_MODULE(CSPdarknet53_tiny);

CSPdarknet53_tinyImpl::CSPdarknet53_tinyImpl() {
   
	conv1 = BasicConv(3, 32, 3, 2);
	conv2 = BasicConv(32, 64, 3, 2);
	resblock_body1 = Resblock_body(64, 64);
	resblock_body2 = Resblock_body(128, 128);
	resblock_body3 = Resblock_body(256, 256);
	conv3 = BasicConv(512, 512, 3);

	register_module("conv1", conv1);
	register_module("conv2", conv2);
	register_module("resblock_body1", resblock_body1);
	register_module("resblock_body2", resblock_body2);
	register_module("resblock_body3", resblock_body3);
	register_module("conv3", conv3);
}

std::vector<torch::Tensor> CSPdarknet53_tinyImpl::forward(torch::Tensor x) {
   
	// 416, 416, 3 -> 208, 208, 32 -> 104, 104, 64
	x = conv1(x);
	x = conv2(x);

	// 104, 104, 64 -> 52, 52, 128
	x = resblock_body1->forward(x)[0];
	// 52, 52, 128 -> 26, 26, 256
	x = resblock_body2->forward(x)[0];
	// 26, 26, 256->xΪ13, 13, 512
#   //        -> feat1Ϊ26,26,256
	auto res_out = resblock_body3->forward(x