python - 如何编写带有预加载的 caffe python 数据层？

Question

如何在执行其他处理时编写异步数据层来预加载批处理？有一些示例代码吗？谢谢

Answer 1

您可以通过多种方式实现您的目标。我会尝试在这里草拟一个选项。

系统的总体 View 是:您有 n 个 Loader 异步加载数据并提供队列。然后，该层从队列中读取 batch_size 项，并将网络馈送到 forward() 函数中。

import caffe, multiprocessing

class Loader(multiprocessing.Process):
  def __init__(self, outq, *args, **kwargs):
    super(Loader, self).__init__()
    self.daemon = True
    self.outq = outq
    self.start()  # start working

  def run(self):
    while True:  # read and never stop at all!
      try:
        # do your magic here
        # assuming you load x,y pairs
        self.outq.put((x[None, ...], y[None, ...]))  # add singleton "batch" dimension
      except Exception as e:
        # handle errors?
        pass

 class MultiProcessInputLayer(caffe.Layer):
   def setup(self, bottom, top):
     # verify no bottoms, right number of tops etc.
     self.dataQ = multiprocessing.Queue()
     for _ in xrange(n):
       Loader(self.dataQ)  # start n Loaders
     # some other stuff here...

   def reshape(self, bottom, top):
     # reshape the inputs to the right sizes

   def forward(self, bottom, top):
     for i in xrange(batch_size):
       item = self.dataQ.get()
       top[0].data[i, ...] = item[0]
       top[1].data[i, ...] = item[1]

   def backward(self, top, propagate_down, bottom):
     pass  # no backward for data layer

我通过艰辛的方式学到的一些提示和技巧:
1. 使用 multiprocessing 而不是 threading 包因为 GIL .
2. 有时(例如，如果 batch_size 非常大)forward() 需要很长时间才能从队列中逐项读取以形成每个批处理。在这种情况下，您可以添加另一个 multiprocessing.Process，它将异步读取 self.dataQ 中的 batch_size 项并将整个批处理写入 self.batchQ。然后 forward() 将在每次调用时仅等待来自 self.batchQ 的单个项。
3. 注意不要复制太多的数据。使用大图像/标签会使所有这些复制成为瓶颈。