针对单机多卡环境的Retinanet目标检测算法实现(简单,明了,易用,中文注释)

牙疼 • 479 次浏览 • 0 个回复 • 2019年10月29日

一. 项目简介

该项目是在SSD-Pytorch项目基础上修改而来.

相关文章见:针对单机多卡环境的SSD目标检测算法实现(Single Shot MultiBox Detector)(简单,明了,易用,中文注释)

完整项目地址:Github,欢迎star, fork.

1.1 项目特点

1.结构化实现.模型各部件均分离实现,模型结构更为清晰明了,便于理解以及后续修改使用.

2.针对单机多卡环境实现.简单直观的多GPU调用实现. 只需通过列表指明需调用的GPU序号即可实现多GPU环境下的训练与测试.如: train_devices=[0,1,3]即可调用第1,2,4三块GPU用于训练.

3.全中文注释.每部分均有详细的中文注释,是你学习Retinanet,加深理解的不二之选.

4.项目结构具有非常高的可移植性,本项目是在SSD-Pytorch项目基础上修改而来,只修改了极少部分代码,可以大量减少重复性的工作.

1.2 Requirements

(均可使用pip进行安装)

1.pytorch

2.opencv-python

3.torchvision >= 0.3.0

4.Vizer

5.visdom

1.3 项目文件结构

文件夹 文件 说明
Data 数据相关
Dataloader 数据加载器类'Our_Dataloader', 'Our_Dataloader_test'
Dataset_VOC VOC格式数据集类
Transfroms 数据Transfroms
Transfroms_tuils Transfroms子方法
Model 模型相关
base_models/Resnet 支持resnet18,34,50,101,152
structs/Anchors retinanet默认检测框生成器
structs/MutiBoxLoss 损失函数
structs/Focal_Loss focal_loss损失函数
structs/Fpn 特征金字塔结构
structs/PostProcess 后处理
structs/Predictor 分类及回归网络
evaler 验证器,用于在数据集上对模型进行验证(测试),计算ap,map
retainnet Retinanet模型类
trainer 训练器,用于在数据集上训练模型
Utils 各种工具
boxs_op 各种框体操作,编码解码,IOU计算,框体格式转换等
Weights 模型权重存放处
pretrained 预训练模型权重存放处,本项目模型并没有训练完毕,因而没有上传训练好的模型,但是训练过程已经过验证
trained 训练过程中默认模型存放处
---- Configs.py 配置文件,包含了模型定义,数据以及训练过程,测试过程等的全部参数,建议备份一份再进行修改
---- Demo_train.py 模型训练的例子,训练过程中的模型会保存在Weights/Our/
---- Demo_eval.py 模型测试的例子,计算模型ap,map
---- Demo_detect_one_image.py 检测单张图片例子
---- Demo_detect_video.py 视频检测例子,传入一个视频,进行检测

二. 例子

本项目Demo均经过测试,可直接运行.

2.1 训练train

针对单机多卡环境的SSD目标检测算法实现(Single Shot MultiBox Detector)(简单,明了,易用,中文注释)一样,项目使用visdom进行训练过程可视化.在运行前请安装并运行visdom.

同样的,训练过程也只支持单机单卡或单机多卡环境,不支持cpu训练.


# -*- coding: utf-8 -*-
# @Author  : LG

from Model import RetainNet, Trainer
from Data import vocdataset
from Configs import _C as cfg
from Data import transfrom,targettransform

# 训练数据集,VOC格式数据集, 训练数据取自 ImageSets/Main/train.txt'
train_dataset=vocdataset(cfg, is_train=True, transform=transfrom(cfg,is_train=True),
                         target_transform=targettransform(cfg))

# 测试数据集,VOC格式数据集, 测试数据取自 ImageSets/Main/eval.txt'
test_dataset = vocdataset(cfg=cfg, is_train=False,
                          transform=transfrom(cfg=cfg, is_train=False),
                          target_transform=targettransform(cfg))

if __name__ == '__main__':
    """
    使用时,请先打开visdom

    命令行 输入  pip install visdom          进行安装 
    输入        python -m visdom.server'    启动
    """

    # 首次调用会下载resnet预训练模型

    # 实例化模型. 模型的具体各种参数在Config文件中进行配置
    net = RetainNet(cfg)
    # 将模型移动到gpu上,cfg.DEVICE.MAINDEVICE定义了模型所使用的主GPU
    net.to(cfg.DEVICE.MAINDEVICE)
    # 初始化训练器,训练器参数通过cfg进行配置;也可传入参数进行配置,但不建议
    trainer = Trainer(cfg)
    # 训练器开始在 数据集上训练模型
    trainer(net, train_dataset)

2.2 验证eval

验证过程支持单机多卡,单机单卡,不支持cpu.

# -*- coding: utf-8 -*-
# @Author  : LG

from Model import RetainNet, Evaler
from Data import vocdataset
from Configs import _C as cfg
from Data import transfrom,targettransform

# 训练数据集,VOC格式数据集, 训练数据取自 ImageSets/Main/train.txt'
train_dataset=vocdataset(cfg, is_train=True, transform=transfrom(cfg,is_train=True),
                         target_transform=targettransform(cfg))

# 测试数据集,VOC格式数据集, 测试数据取自 ImageSets/Main/eval.txt'
test_dataset = vocdataset(cfg=cfg, is_train=False,
                          transform=transfrom(cfg=cfg, is_train=False),
                          target_transform=targettransform(cfg))

if __name__ == '__main__':
    # 模型测试只支持GPU单卡或多卡,不支持cpu
    net = RetainNet(cfg)
    # 将模型移动到gpu上,cfg.DEVICE.MAINDEVICE定义了模型所使用的主GPU
    net.to(cfg.DEVICE.MAINDEVICE)
    # 模型从权重文件中加载权重
    net.load_pretrained_weight('XXX.pkl')
    # 初始化验证器,验证器参数通过cfg进行配置;也可传入参数进行配置,但不建议
    evaler = Evaler(cfg, eval_devices=None)
    # 验证器开始在数据集上验证模型
    ap, map = evaler(model=net,
                     test_dataset=test_dataset)
    print(ap)
    print(map)

2.3 检测Detect

单次检测过程支持单机单卡,cpu.

2.3.1 单张图片检测

# -*- coding: utf-8 -*-
# @Author  : LG
from Model import RetainNet
from Configs import _C as cfg
from PIL import Image
from matplotlib import pyplot as plt

# 实例化模型
net = RetainNet(cfg)
# 使用cpu或gpu
net.to('cuda')
# 模型从权重文件中加载权重
net.load_pretrained_weight('XXX.pkl')
# 打开图片
image = Image.open("XXX.jpg")
# 进行检测, 分别返回 绘制了检测框的图片数据/回归框/标签/分数.
drawn_image, boxes, labels, scores = net.Detect_single_img(image=image,score_threshold=0.5)

plt.imsave('XXX_det.jpg',drawn_image)
plt.imshow(drawn_image)
plt.show()

2.3.2 视频检测

# -*- coding: utf-8 -*-
# @Author  : LG
from Model import RetainNet
from Configs import _C as cfg

# 实例化模型
net = RetainNet(cfg)
# 使用cpu或gpu
net.to('cuda')
# 模型从权重文件中加载权重
net.load_pretrained_weight('XXX.pkl')

video_path = 'XXX.mp4'

# 进行检测,
# if save_video_path不为None,则不保存视频,如需保存视频save_video_path=XXX.mp4 ,
# show=True,实时显示检测结果
net.Detect_video(video_path=video_path, score_threshold=0.02, save_video_path=None, show=True)

说明

该项目是在SSD-Pytorch项目基础上修改而来.

由于部分原因,不会上传预训练模型,但是训练过程已经过验证,代码无误.

完整项目地址:Github,欢迎star, fork.

仅限用于交流学习,如需引用,请联系作者.


原创文章,转载请注明 :针对单机多卡环境的Retinanet目标检测算法实现(简单,明了,易用,中文注释) - pytorch中文网
原文出处: https://ptorch.com/news/254.html
问题交流群 :168117787
提交评论
要回复文章请先登录注册
用户评论
  • 没有评论
Pytorch是什么?关于Pytorch!