Class BatchDataByShape

tensorpack/dataflow/common.py:193–229

Group datapoints of the same shape together to batches. It doesn't require input DataFlow to be homogeneous anymore: it can have datapoints of different shape, and batches will be formed from those who have the same shape. Note: It is implemented by a dict{shape -> data

class BatchDataByShape(BatchData):
    """
    Group datapoints of the same shape together to batches.
    It doesn't require input DataFlow to be homogeneous anymore: it can have
    datapoints of different shape, and batches will be formed from those who
    have the same shape.

    Note:
        It is implemented by a dict{shape -> datapoints}.
        Therefore, datapoints of uncommon shapes may never be enough to form a batch and
        never get generated.
    """
    def __init__(self, ds, batch_size, idx):
        """
        Args:
            ds (DataFlow): input DataFlow. ``dp[idx]`` has to be an :class:`np.ndarray`.
            batch_size (int): batch size
            idx (int): ``dp[idx].shape`` will be used to group datapoints.
                Other components are assumed to be batch-able.
        """
        super(BatchDataByShape, self).__init__(ds, batch_size, remainder=False)
        self.idx = idx

    def reset_state(self):
        super(BatchDataByShape, self).reset_state()
        self.holder = defaultdict(list)
        self._guard = DataFlowReentrantGuard()

    def __iter__(self):
        with self._guard:
            for dp in self.ds:
                shp = dp[self.idx].shape
                holder = self.holder[shp]
                holder.append(dp)
                if len(holder) == self.batch_size:
                    yield BatchData.aggregate_batch(holder)
                    del holder[:]


class FixedSizeData(ProxyDataFlow):