工程实现 || YOLOv3(3) --- 数据读取的多进程改造 dataset.py

前言：

yolo系列的论文阅读
论文阅读 || 深度学习之目标检测 重磅出击YOLOv3
论文阅读 || 深度学习之目标检测yolov2
论文阅读 || 深度学习之目标检测yolov1
 
该篇讲解的工程连接是：
tensorflow的yolov3：https://github.com/YunYang1994/tensorflow-yolov3
 
自己对该工程的解析博客：
工程实现 || YOLOv3(1) —使用自己的数据集训练yolov3
工程实现 || YOLOv3(2) —dataset.py解析
工程实现 || YOLOv3(3) —dataset.py的多进程改造
工程实现 || YOLOv(4) —yolov3.py 网络结构的搭建和loss的定义
工程实现 || YOLOv3(5) —train.py
工程实现 || YOLOv3(6) — anchorboxes的获取 kmeans.py
 
数据读取部分，需要经过较多处理，CPU的计算经常会占用大量时间。导致GPU经常处于等待的状态，GPU利用率较低，整个训练时间大大延长。 所以我们希望快速读取数据，来提高GPU的利用率。
该篇是对数据读取部分做多进程的改造，利用python的多进程，来完成数据的快速读取。
其实 tensorflow有相关函数来完成多进程，但个人感觉使用起来并不简便
 

from multiprocessing import Queue, Process

1 原工程的数据读取的回顾

上一篇链接中也提到

class Dataset(object):
  """implement Dataset here"""
  def __init__(self, dataset_type):
      ...
  def __iter__(self):
      return self
  def __next__(self):
      ...
  def __len__(self):
     return self.num_batchs

---

其中主要内容为def __next__(self)，从多进程改造的角度来看，该函数分为3部分

• 创建存放【输入图片】和【label 】的数组
• 设置读取数量的判定条件 || if / for / while相关内容
• 数据增强处理和label的制作，并return

 
具体定义为

    def __next__(self):

       with tf.device('/cpu:0'):

           """part 1：创建存放【输入图片】和【label 】的数组"""
           self.train_input_size = random.choice(self.train_input_sizes)
           self.train_output_sizes = self.train_input_size // self.strides
           # 创建存放【输入图片】的数组
           batch_image = np.zeros((self.batch_size, self.train_input_size, self.train_input_size, 3))
           # 创建存放【label】的数组
           batch_label_sbbox = np.zeros((self.batch_size, self.train_output_sizes[0], self.train_output_sizes[0], self.anchor_per_scale, 5 + self.num_classes))
           batch_label_mbbox = np.zeros((self.batch_size, self.train_output_sizes[1], self.train_output_sizes[1], self.anchor_per_scale, 5 + self.num_classes))
           batch_label_lbbox = np.zeros((self.batch_size, self.train_output_sizes[2], self.train_output_sizes[2], self.anchor_per_scale, 5 + self.num_classes))
            # 创建存放 在3个尺度下，负责预测的bboxes
           batch_sbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))
           batch_mbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))
           batch_lbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))

           """part 2：设置读取数量的判定条件 || if / for / while相关内容"""
           num = 0       # 批内的计数器
           if self.batch_count < self.num_batchs:     # 当【已读取的批数】小于【一轮总批数】
               while num < self.batch_size:           # 当【批内读取个数】小于【batch】
                   index = self.batch_count * self.batch_size + num  # 【index】为一轮内已经读取的数据个数
                   if index >= self.num_samples: index -= self.num_samples  # 如果【index】大于【数据总量】，将index置0
                   
                   """part 3： 数据增强处理和label的制作，并return """
                   # 读取训练集或验证集的数据信息（输入图片路径、bboxes）
                   annotation = self.annotations[index]
                   # 数据增强：随机翻转、随机裁剪、随机平移、将图片缩放和填充到target_shape（相同规则处理bboxes）
                   image, bboxes = self.parse_annotation(annotation)
                   # 结合anchorbox等信息，将bboxes的信息，转化为神经网络训练的label
                   label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes = self.preprocess_true_boxes(bboxes)

                   batch_image[num, :, :, :] = image
                   batch_label_sbbox[num, :, :, :, :] = label_sbbox
                   batch_label_mbbox[num, :, :, :, :] = label_mbbox
                   batch_label_lbbox[num, :, :, :, :] = label_lbbox
                   batch_sbboxes[num, :, :] = sbboxes
                   batch_mbboxes[num, :, :] = mbboxes
                   batch_lbboxes[num, :, :] = lbboxes
                   num += 1
               self.batch_count += 1
               return batch_image, batch_label_sbbox, batch_label_mbbox, batch_label_lbbox, \
                      batch_sbboxes, batch_mbboxes, batch_lbboxes
           else:
               self.batch_count = 0
               np.random.shuffle(self.annotations)
               raise StopIteration

2 工程改造

为了尽量小的变动源代码，这里重新创建个脚本load_data.py，来实现多进程。
 
思路： 使用队列和多进程对于一个dataset的对象，属性增加两个队列。
多进程获取数据信息（图片路径和bboxes）的来源与Q_name，多进程处理后的数据放入Q_data。

• Q_data： 存放的是神经网络训练时的label。创建多个进程进行数据读取，然后都放入Q_data对列中，送入神经网络的placehoder然后进行训练。
• Q_name： 存放文本读取的数据信息（图片路径和bboxes）。
   

疑问： 为什么要创建了读取数据名字等信息的队列呢？
因为按照多进程读取数据时，每个进程读取自己的数据名字的列表，相同的数据可能在一定范围内多次出现。这样就相当于数据集拷贝了n（多进程数量）倍，然后进行乱序读取（虽然这种方式对训练的结果影响不是很大）。所以创建个名字的对列，多进程获取数据信息的来源与Q_name，多进程处理后的数据放入Q_data。

---

2.1 dataset.py的改动

2.1.1 def __init__()：部分的改动

该部分主要是添加了类Dataset的属性Q_data、Q_name。以便其他脚本对类Dataset的这两个属性操作的有效性。

class Dataset(object):
   def __init__(self, dataset_type):
        self.annot_path  = cfg.TRAIN.ANNOT_PATH if dataset_type == 'train' else cfg.TEST.ANNOT_PATH
        self.input_sizes = cfg.TRAIN.INPUT_SIZE if dataset_type == 'train' else cfg.TEST.INPUT_SIZE
        self.batch_size  = cfg.TRAIN.BATCH_SIZE if dataset_type == 'train' else  cfg.TEST.BATCH_SIZE
        self.data_aug    = cfg.TRAIN.DATA_AUG   if dataset_type == 'train' else cfg.TEST.DATA_AUG
        self.train_input_sizes = cfg.TRAIN.INPUT_SIZE
        self.strides = np.array(cfg.YOLO.STRIDES)
        self.classes = utils.read_class_names(cfg.YOLO.CLASSES)
        self.num_classes = len(self.classes)
        self.anchors = np.array(utils.get_anchors(cfg.YOLO.ANCHORS))
        self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE
        self.max_bbox_per_scale = 150
           
        self.annotations = self.load_annotations(dataset_type)
        self.num_samples = len(self.annotations)    # 样本的数量
        self.num_batchs = int(np.ceil(self.num_samples / self.batch_size))  # 一轮读取批
        self.batch_count = 0
           
         """添加的部分"""
        self.num_trianepoch = cfg.TRAIN.FISRT_STAGE_EPOCHS + cfg.TRAIN.SECOND_STAGE_EPOCHS
        self.Q_name = []      
        self.Q_data = []

 

---

2.1.2 def __next__()：部分的改动

前面提到过：def __next__(self)，从多进程改造的角度来看，该函数分为3部分

• 创建存放【输入图片】和【label 】的数组
• 设置读取数量的判定条件 || if / for / while相关内容
• 数据增强处理和label的制作，并return

 
在def __next__(self)中，保留第一、三部分的操作。

   def Prepare(self):
     with tf.device('/cpu:0'):
       self.train_input_size = random.choice(self.train_input_sizes)
       self.train_output_sizes = self.train_input_size // self.strides
       # 创建存放【输入图片】的数组
       batch_image = np.zeros((self.batch_size, self.train_input_size, self.train_input_size, 3))
       # 创建存放【label】的数组
       batch_label_sbbox = np.zeros((self.batch_size, self.train_output_sizes[0], self.train_output_sizes[0], self.anchor_per_scale, 5 + self.num_classes))
       batch_label_mbbox = np.zeros((self.batch_size, self.train_output_sizes[1], self.train_output_sizes[1], self.anchor_per_scale, 5 + self.num_classes))
       batch_label_lbbox = np.zeros((self.batch_size, self.train_output_sizes[2], self.train_output_sizes[2], self.anchor_per_scale, 5 + self.num_classes))
        # 创建存放 在3个尺度下，负责预测的bboxes
       batch_sbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))
       batch_mbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))
       batch_lbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))


  def readdata(self, annotation, num):
    # 读取训练集或验证集的数据信息（输入图片路径、bboxes）
    # annotation = self.annotations[index]

    # 数据增强：随机翻转、随机裁剪、随机平移、将图片缩放和填充到target_shape（相同规则处理bboxes）
    # 此时bboxes为【左上角右下角】形式
    image, bboxes = self.parse_annotation(annotation)
    # 结合anchorbox等信息，将bboxes的信息，转化为神经网络训练的label
    label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes = self.preprocess_true_boxes(bboxes)

    self.batch_image[num, :, :, :] = image
    self.batch_label_sbbox[num, :, :, :, :] = label_sbbox
    self.batch_label_mbbox[num, :, :, :, :] = label_mbbox
    self.batch_label_lbbox[num, :, :, :, :] = label_lbbox
    self.batch_sbboxes[num, :, :] = sbboxes
    self.batch_mbboxes[num, :, :] = mbboxes
    self.batch_lbboxes[num, :, :] = lbboxes

---

2.2 load_data.py的编写

2.2.1 class QueueGene()。

实现的功能：将数据存到队列中。包含原本工程中数据读取中的def __next__(self)的第二部分，设置读取数量的判定条件 || if / for / while相关内容。

from multiprocessing import Queue, Process
from core.dataset import *
import random


class QueueGene():
   def __init__(self):
       print("创建数据队列初始化成功")
     
   def getname(self, datagene):
       print("向队列Q_name中，存放文本读取的数据信息（图片路径和bboxes）")
       
       for i in range(datagene.num_trianepoch):
           random.shuffle(datagene.annotations)  #每一轮的数据信息进行打乱
           for j in range(datagene.num_samples):
               datagene.Q_name.put(datagene.annotations[j])  # 数据逐个放入Q_name中
               
   def Data(self, datagene, thread):
       print("向队列Q_data中，存放神经网络训练时的label")
      
       datagene.Prepare()

       name = []
       while 1:
           if datagene.batch_count < datagene.num_batchs:   # 当【读取了几批】小于【一轮总批数】
               num = 0                             # 统计批内读取个数
               while num < datagene.batch_size:        #【批内读取数据个数】小于【一个batch数值】
                   namefile = datagene.Q_name.get()    ##  获取数据名字等信息
                   datagene.readdata(namefile, num)   ## 读取并处理数据 
                   name.append(namefile)
                   num += 1
               datagene.batch_count += 1 # 统计一轮的训练，读取了几个批次
               datagene.Q_data.put([name, thread,     ## 存放name、thread是为了自己测试数据读取是否正确所用
                           datagene.batch_image, datagene.batch_label_sbbox,
                           datagene.batch_label_mbbox, datagene.batch_label_lbbox, datagene.batch_sbboxes,
                           datagene.batch_mbboxes, datagene.batch_lbboxes])
               name = []
           else:
               datagene.batch_count = 0

---

2.2.2 def load_data_start()

前面实现了class QueueGene()，接下来就要为这个类创建个实例，来运行相关功能。

def load_data_start(Q_datanum1=1000, Q_datanum2=1000, P1=4, P2=2):

  traindatagene = Dataset("train")
  testdatagene = Dataset("test")
  traindatagene.Q_name = Queue(traindatagene.num_samples)  # 读取训练数据的名字的队列
  testdatagene.Q_name = Queue(testdatagene.num_samples)  # 读取验证数据的名字的队列
  traindatagene.Q_data = Queue(Q_datanum1)  # 读取训练数据的队列
  testdatagene.Q_data = Queue(Q_datanum2)  # 读取验证数据的队列

  queuegene = QueueGene()
  Process(target=queuegene.getname, args=(traindatagene,)).start()  # 创建进程，读取训练数据的名字等信息
  Process(target=queuegene.getname, args=(testdatagene,)).start()  # 创建进程，读取验证数据的名字等信息
  for thread in range(P1):
      Process(target=queuegene.Data, args=(traindatagene, thread)).start()  # 创建多进程，读取训练数据
  for thread in range(P2):
      Process(target=queuegene.Data, args=(testdatagene, thread)).start())  # 创建多进程，读取验证数据

  return traindatagene.Q_data, testdatagene.Q_data
  
"""调用函数，用于测试函数获取数据 实现的是否有误"""
if __name__ == '__main__':
  Q_traindata, Q_validdata = load_data_start(100,0, 1,0)
  for i in range(10):
      A = Q_traindata.get()
      

有任何问题，欢迎讨论