batch_manager.py 文件源码

def __init__(self, source, batch_size, shuffle=True, use_prefetch=True, capacity=32):
        ## open LMDB
        self.env = lmdb.open(source, readonly=True)
        self.txn = self.env.begin()
        self.cur = self.txn.cursor()

        self.batch_size = batch_size
        self.shuffle = shuffle
        self.use_prefetch = use_prefetch
        self.capacity = capacity
        self.num_data = int(self.txn.stat()['entries'])

        self.reset_inds()

        if self.use_prefetch:
            self.batch_queue = Queue(capacity)
            self.proc = Process(target=self._worker)
            self.proc.start()
            def cleanup():
                self.proc.terminate()
                self.proc.join()
            import atexit
            atexit.register(cleanup)