我参加 NVIDIA Sky Hackathon---yolov5数据标注+模型训练

最新推荐文章于 2024-07-09 20:33:07 发布

碳酸碳

最新推荐文章于 2024-07-09 20:33:07 发布

阅读量910

点赞数

分类专栏： NVIDIA黑客松文章标签：深度学习计算机视觉 python

本文链接：https://blog.csdn.net/qq_41679000/article/details/128078428

版权

NVIDIA黑客松专栏收录该内容

3 篇文章 0 订阅

订阅专栏

该博客介绍了在图像分类任务中使用LabelImg和CVAT进行高效标注的过程，以及如何将XML标签转换为YOLO格式。小组通过在图片上添加椒盐噪声增强模型鲁棒性，并利用RTX3060 GPU及远程服务器进行YOLOv5模型训练。在训练优化阶段，调整了超参数如马赛克功能和图像变换，最终模型达到了预期的准确率。

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

在标注过程中，将图片分为三类，分别为banana、zhiban、yilaguan三个类别，在yolov5的数据labels标注种分别将其划为0，1，2三类。在确定了分类规则后，使用labelimg进行标注。在此附上labelimg的项目代码：https://github.com/tzutalin/labelImg

标注工作完成后，因本项目采用的是yolo格式的数据集，所以在模型导出的过程中需要设定为yolo格式。小组同时就官方给的xml类型的标注labels编写脚本进行转换，待所有数据格式转换完毕后，将labels中的框画到图片中以测试效果：

在此小组也尝试了一种更为高效的标注平台CVAT，此平台需要在linux平台中部署，并且可以开启一个服务器，小组成员所有人都可以同时在这个标注平台上进行标注，管理员可以在后台给每个成员创建账号，并在其中给每个成员创建标注任务，就并发性和界面UI的设计上来说要优于labelmg不少。同样，在此附上CVAT的项目代码：GitHub - opencv/cvat: Annotate better with CVAT, the industry-leading data engine for machine learning. Used and trusted by teams at any scale, for data of any scale.

其次，我们针对模型的鲁棒性进行优化，对于图片大小编写脚本对图片随机添加3000个大小为3*3像素点的椒盐噪声，虽然会相对降低模型的精确度，但是就综合性能来说还是有一定提升。

核心代码如下：

for i in range(0, temp):

if random.randint(1, temp) % 2 == 0:

newheight= random.randint(0, height)

newwidth=random.randint(0, width)

img[newheight:min(newheight+3,height), newwidth:min(newwidth+3,width)] = (255, 255, 255)

if random.randint(1, temp) % 2 != 0:

newheight= random.randint(0, height)

newwidth=random.randint(0, width)

img[newheight:min(newheight+3,height), newwidth:min(newwidth+3,width)] = (0, 0, 0)

模型训练

硬件选择

对于模型的训练过程，我们基于小组成员的机器配置情况，考虑到充分利用到GPU的性能，我们使用小组成员的RTX3060显卡，使用yolov5x预训练模型，以batch 8训练300批，在经过效果评估后，虽然效果差强人意，但是为了追求更好的训练效果，突破硬件设备限制，我们分别采用了Kaggle上的GPU训练和租用阿里的服务器分别进行训练。经过13小时的训练，训练完成并获取到了best.pt和last.pt