Method _enqueue_batch

caffe2/python/data_workers.py:259–316

This pulls data from the python-side queue and collects them into batch-sized pieces, unless dont_rebatch is set to true.

(self, data_input_coordinator)

                self._enqueue(b, q, c)

    def _enqueue_batch(self, data_input_coordinator):
        '''
        This pulls data from the python-side queue and collects them
        into batch-sized pieces, unless dont_rebatch is set to true.
        '''
        if self._dont_rebatch:
            self._enqueue_batch_direct(data_input_coordinator)
            return

        cur_batch = [np.array([]) for d in self._input_blob_names]
        first_batch_col = self._batch_columns[0]

        # Collect data until we have a full batch size
        while (
            cur_batch[0].shape[0] == 0 or
            cur_batch[0].shape[first_batch_col] < self._batch_size
        ) and data_input_coordinator.is_active():
            chunk = self._get(data_input_coordinator)
            if chunk is None:
                continue

            for j, chunk_elem in enumerate(chunk):
                if cur_batch[j].shape[0] == 0:
                    cur_batch[j] = chunk_elem.copy()
                else:
                    cur_batch[j] = np.append(
                        cur_batch[j], chunk_elem, axis=self._batch_columns[j]
                    )

        start_time = time.time()
        try:
            # Return data over the batch size back to queue
            if cur_batch[0].shape[0] > 0 and cur_batch[0].shape[
                first_batch_col
            ] > self._batch_size:
                leftover = []
                trimmed_batch = []
                for j, b in enumerate(cur_batch):
                    [c, l] = np.split(
                        b, [self._batch_size], axis=self._batch_columns[j]
                    )
                    leftover.append(l)
                    trimmed_batch.append(c)
                cur_batch = trimmed_batch
                try:
                    self._internal_queue.put(leftover, block=False)
                except Queue.Full:
                    pass

                assert cur_batch[0].shape[first_batch_col] == self._batch_size

            if data_input_coordinator.is_active():
                for b, q, c in zip(
                    self._input_blob_names, self._queues, cur_batch
                ):
                    self._enqueue(b, q, c)
        finally:
            self._metrics.put_metric('enqueue_time', time.time() - start_time)