Source code for gluonnlp.data.batchify.batchify

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Batchify functions. They can be used in Gluon data loader to help combine individual samples
into batches for fast processing."""
__all__ = ['Stack', 'Pad', 'Tuple', 'List', 'NamedTuple', 'Dict']

import warnings
import math
from typing import (Dict as t_Dict, Callable as t_Callable, List as t_List, Tuple as t_Tuple,
                    AnyStr, Union as t_Union)

import numpy as np
import mxnet as mx


def _pad_arrs_to_max_length(arrs, pad_axis, pad_val, use_shared_mem, dtype, round_to=None):
    """Inner Implementation of the Pad batchify

    Parameters
    ----------
    arrs : list
    pad_axis : int
    pad_val : number
    use_shared_mem : bool, default False

    Returns
    -------
    ret : NDArray
    original_length : NDArray
    """
    if isinstance(arrs[0], mx.nd.NDArray):
        dtype = arrs[0].dtype if dtype is None else dtype
        arrs = [arr.asnumpy() for arr in arrs]
    elif not isinstance(arrs[0], np.ndarray):
        arrs = [np.asarray(ele) for ele in arrs]
    else:
        dtype = arrs[0].dtype if dtype is None else dtype

    original_length = [ele.shape[pad_axis] for ele in arrs]
    max_size = max(original_length)
    if round_to is not None:
        max_size = round_to * math.ceil(max_size / round_to)

    ret_shape = list(arrs[0].shape)
    ret_shape[pad_axis] = max_size
    ret_shape = (len(arrs), ) + tuple(ret_shape)

    ret = np.full(shape=ret_shape, fill_value=pad_val, dtype=dtype)

    for i, arr in enumerate(arrs):
        if arr.shape[pad_axis] == max_size:
            ret[i] = arr
        else:
            slices = [slice(None) for _ in range(arr.ndim)]
            slices[pad_axis] = slice(0, arr.shape[pad_axis])
            if slices[pad_axis].start != slices[pad_axis].stop:
                slices = [slice(i, i + 1)] + slices
                ret[tuple(slices)] = arr

    ctx = mx.Context('cpu_shared', 0) if use_shared_mem else mx.cpu()
    ret = mx.nd.array(ret, ctx=ctx, dtype=dtype)
    original_length = mx.nd.array(original_length, ctx=ctx, dtype=np.int32)

    return ret, original_length


def _stack_arrs(arrs, use_shared_mem, dtype):
    if isinstance(arrs[0], mx.nd.NDArray):
        dtype = arrs[0].dtype if dtype is None else dtype
        if use_shared_mem:
            out = mx.nd.empty((len(arrs),) + arrs[0].shape, dtype=dtype,
                              ctx=mx.Context('cpu_shared', 0))
            return mx.nd.stack(*arrs, out=out)
        else:
            return mx.nd.stack(*arrs)
    else:
        out = np.asarray(arrs)
        dtype = out.dtype if dtype is None else dtype
        if use_shared_mem:
            return mx.nd.array(out, ctx=mx.Context('cpu_shared', 0), dtype=dtype)
        else:
            return mx.nd.array(out, dtype=dtype)


[docs]class Stack: r"""Stack the input data samples to construct the batch. The N input samples must have the same shape/length and will be stacked to construct a batch. Parameters ---------- dtype : str or numpy.dtype, default None The value type of the output. If it is set to None, the input data type is used. Examples -------- >>> import gluonnlp.data.batchify as bf >>> # Stack multiple lists >>> a = [1, 2, 3, 4] >>> b = [4, 5, 6, 8] >>> c = [8, 9, 1, 2] >>> bf.Stack()([a, b, c]) <BLANKLINE> [[1 2 3 4] [4 5 6 8] [8 9 1 2]] <NDArray 3x4 @cpu_shared(0)> >>> # Stack multiple numpy.ndarrays >>> a = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]) >>> b = np.array([[5, 6, 7, 8], [1, 2, 3, 4]]) >>> bf.Stack()([a, b]) <BLANKLINE> [[[1 2 3 4] [5 6 7 8]] <BLANKLINE> [[5 6 7 8] [1 2 3 4]]] <NDArray 2x2x4 @cpu_shared(0)> >>> # Stack multiple NDArrays >>> a = mx.nd.array([[1, 2, 3, 4], [5, 6, 7, 8]]) >>> b = mx.nd.array([[5, 6, 7, 8], [1, 2, 3, 4]]) >>> bf.Stack()([a, b]) <BLANKLINE> [[[1. 2. 3. 4.] [5. 6. 7. 8.]] <BLANKLINE> [[5. 6. 7. 8.] [1. 2. 3. 4.]]] <NDArray 2x2x4 @cpu_shared(0)> """ def __init__(self, dtype=None): self._dtype = dtype
[docs] def __call__(self, data): """Batchify the input data Parameters ---------- data : list The input data samples Returns ------- batch_data : NDArray """ return _stack_arrs(data, True, self._dtype)
[docs]class Pad: """Return a callable that pads and stacks data. Parameters ---------- axis : int, default 0 The axis to pad the arrays. The arrays will be padded to the largest dimension at `axis`. For example, assume the input arrays have shape (10, 8, 5), (6, 8, 5), (3, 8, 5) and the `axis` is 0. Each input will be padded into (10, 8, 5) and then stacked to form the final output, which has shape´╝ł3, 10, 8, 5). pad_val : float or int, default 0 The padding value. ret_length : bool, default False Whether to return the valid length in the output. dtype : str or numpy.dtype, default None The value type of the output. If it is set to None, the input data type is used. round_to : int, default None If specified, the padded dimension will be rounded to be multiple of this argument. Examples -------- >>> import gluonnlp.data.batchify as bf >>> # Inputs are multiple lists >>> a = [1, 2, 3, 4] >>> b = [4, 5, 6] >>> c = [8, 2] >>> bf.Pad(pad_val=0)([a, b, c]) <BLANKLINE> [[1. 2. 3. 4.] [4. 5. 6. 0.] [8. 2. 0. 0.]] <NDArray 3x4 @cpu_shared(0)> >>> # Also output the lengths >>> a = [1, 2, 3, 4] >>> b = [4, 5, 6] >>> c = [8, 2] >>> batch, length = bf.Pad(pad_val=0, ret_length=True)([a, b, c]) >>> batch <BLANKLINE> [[1. 2. 3. 4.] [4. 5. 6. 0.] [8. 2. 0. 0.]] <NDArray 3x4 @cpu_shared(0)> >>> length <BLANKLINE> [4 3 2] <NDArray 3 @cpu_shared(0)> >>> # Inputs are multiple ndarrays >>> a = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]) >>> b = np.array([[5, 8], [1, 2]]) >>> bf.Pad(axis=1, pad_val=-1)([a, b]) <BLANKLINE> [[[ 1 2 3 4] [ 5 6 7 8]] <BLANKLINE> [[ 5 8 -1 -1] [ 1 2 -1 -1]]] <NDArray 2x2x4 @cpu_shared(0)> """ def __init__(self, axis=0, pad_val=None, ret_length=False, dtype=None, round_to=None): self._axis = axis assert isinstance(axis, int), 'axis must be an integer! ' \ 'Received axis=%s, type=%s.' % (str(axis), str(type(axis))) self._pad_val = 0 if pad_val is None else pad_val self._ret_length = ret_length self._dtype = dtype self._warned = False self._round_to = round_to if pad_val is None: warnings.warn( 'Padding value is not given and will be set automatically to 0 ' 'in data.batchify.Pad(). ' 'Please check whether this is intended ' '(e.g. value of padding index in the vocabulary).')
[docs] def __call__(self, data): """Batchify the input data. The input can be list of numpy.ndarray, list of numbers or list of mxnet.nd.NDArray. Inputting mxnet.nd.NDArray is discouraged as each array need to be converted to numpy for efficient padding. The arrays will be padded to the largest dimension at `axis` and then stacked to form the final output. In addition, the function will output the original dimensions at the `axis` if ret_length is turned on. Parameters ---------- data : List[np.ndarray] or List[List[dtype]] or List[mx.nd.NDArray] List of samples to pad and stack. Returns ------- batch_data: NDArray Data in the minibatch. Shape is (N, ...) valid_length: NDArray, optional The sequences' original lengths at the padded axis. Shape is (N,). This will only be returned in `ret_length` is True. """ if isinstance(data[0], mx.nd.NDArray) and not self._warned: self._warned = True warnings.warn( 'Using Pad with NDArrays is discouraged for speed reasons. ' 'Instead you should pad your data while it is still a list ' 'and before converting to an NDArray. ' 'Alternatively you can consider inputting a numpy.ndarray.') if isinstance(data[0], (mx.nd.NDArray, np.ndarray, list)): padded_arr, original_length = _pad_arrs_to_max_length(data, self._axis, self._pad_val, True, self._dtype, round_to=self._round_to) if self._ret_length: return padded_arr, original_length else: return padded_arr else: raise NotImplementedError
[docs]class Tuple: """Wrap multiple batchify functions together. The input functions will be applied to the corresponding input fields. Each data sample should be a list or tuple containing multiple attributes. The `i`th batchify function stored in `Tuple` will be applied on the `i`th attribute. For example, each data sample is (nd_data, label). You can wrap two batchify functions using `Tuple(DataBatchify, LabelBatchify)` to batchify nd_data and label correspondingly. Parameters ---------- fn : list or tuple or callable The batchify functions to wrap. *args : tuple of callable The additional batchify functions to wrap. Examples -------- >>> import gluonnlp.data.batchify as bf >>> a = ([1, 2, 3, 4], 0) >>> b = ([5, 7], 1) >>> c = ([1, 2, 3, 4, 5, 6, 7], 0) >>> f1, f2 = bf.Tuple(bf.Pad(pad_val=0), bf.Stack())([a, b]) >>> f1 <BLANKLINE> [[1. 2. 3. 4.] [5. 7. 0. 0.]] <NDArray 2x4 @cpu_shared(0)> >>> f2 <BLANKLINE> [0 1] <NDArray 2 @cpu_shared(0)> """ def __init__(self, fn, *args): if isinstance(fn, (list, tuple)): assert len(args) == 0, 'Input pattern not understood. The input of Tuple can be ' \ 'Tuple(A, B, C) or Tuple([A, B, C]) or Tuple((A, B, C)). ' \ 'Received fn=%s, args=%s' % (str(fn), str(args)) self._fn = fn else: self._fn = (fn, ) + args for i, ele_fn in enumerate(self._fn): assert hasattr(ele_fn, '__call__'), 'Batchify functions must be callable! ' \ 'type(fn[%d]) = %s' % (i, str(type(ele_fn)))
[docs] def __call__(self, data): """Batchify the input data. Parameters ---------- data : list The samples to batchfy. Each sample should contain N attributes. Returns ------- ret : tuple A tuple of length N. Contains the batchified result of each attribute in the input. """ assert len(data[0]) == len(self._fn),\ 'The number of attributes in each data sample should contains' \ ' {} elements'.format(len(self._fn)) ret = [] for i, ele_fn in enumerate(self._fn): ret.append(ele_fn([ele[i] for ele in data])) return tuple(ret)
[docs]class List: """Simply forward the list of input data. This is particularly useful when the Dataset contains textual data and in conjonction with the `Tuple` batchify function. Examples -------- >>> import gluonnlp.data.batchify as bf >>> a = ([1, 2, 3, 4], "I am using MXNet") >>> b = ([5, 7, 2, 5], "Gluon rocks!") >>> c = ([1, 2, 3, 4], "Batchification!") >>> _, l = bf.Tuple(bf.Stack(), bf.List())([a, b, c]) >>> l ['I am using MXNet', 'Gluon rocks!', 'Batchification!'] """
[docs] def __call__(self, data: t_List) -> t_List: """ Parameters ---------- data The list of samples Returns ------- ret The input list """ return list(data)
[docs]class Dict: """Wrap multiple batchify functions together and apply it to merge inputs from a dict. The generated batch samples are stored as a dict with the same keywords. Each data sample should be a dict and the fn corresponds to `key` will be applied on the input with the keyword `key`. For example, each data sample is {'data': nd_data, 'label': nd_label}. You can merge the data and labels using `Dict({'data': DataBatchify, 'label': LabelBatchify})` to batchify the nd_data and nd_label. Parameters ---------- fn_dict A dictionary that contains the key-->batchify function mapping. Examples -------- >>> from gluonnlp.data.batchify import Dict, Pad, Stack >>> a = {'data': [1, 2, 3, 4], 'label': 0} >>> b = {'data': [5, 7], 'label': 1} >>> c = {'data': [1, 2, 3, 4, 5, 6, 7], 'label': 0} >>> batchify_fn = Dict({'data': Pad(pad_val=0), 'label': Stack()}) >>> sample = batchify_fn([a, b, c]) >>> sample['data'] <BLANKLINE> [[1. 2. 3. 4. 0. 0. 0.] [5. 7. 0. 0. 0. 0. 0.] [1. 2. 3. 4. 5. 6. 7.]] <NDArray 3x7 @cpu_shared(0)> >>> sample['label'] <BLANKLINE> [0 1 0] <NDArray 3 @cpu_shared(0)> """ def __init__(self, fn_dict: t_Dict[AnyStr, t_Callable]): self._fn_dict = fn_dict if not isinstance(fn_dict, dict): raise ValueError('Input must be a dictionary! type of input = {}' .format(type(fn_dict))) for fn in fn_dict.values(): if not hasattr(fn, '__call__'): raise ValueError('Elements of the dictionary must be callable!') self._fn_dict = fn_dict
[docs] def __call__(self, data: t_List[t_Dict]) -> t_Dict: """ Parameters ---------- data The samples to batchify. Each sample should be a dictionary Returns ------- ret The resulting dictionary that stores the merged samples. """ ret = dict() for k, ele_fn in self._fn_dict.items(): ret[k] = ele_fn([ele[k] for ele in data]) return ret
[docs]class NamedTuple: """Wrap multiple batchify functions together and apply it to merge inputs from a namedtuple. The generated batch samples are stored as a namedtuple with the same structure. Each data sample should be a namedtuple. The `i`th batchify function stored in `NamedTuple` will be applied on the `i`th attribute of the namedtuple data. For example, each data sample is Sample(data=nd_data, label=nd_label). You can wrap two batchify functions using `NamedTuple(Sample, {'data': DataBatchify, 'label': LabelBatchify})` to batchify nd_data and nd_label correspondingly. The result will be stored as a Sample object and you can access the data and label via `sample.data` and `sample.label`, correspondingly. Parameters ---------- container : NamedTuple class The object that constructs the NamedTuple. fn_info The information of the inner batchify functions. Examples -------- >>> from gluonnlp.data.batchify import NamedTuple, Pad, Stack >>> from collections import namedtuple >>> SampleData = namedtuple('SampleData', ['data', 'label']) >>> a = SampleData([1, 2, 3, 4], 0) >>> b = SampleData([5, 7], 1) >>> c = SampleData([1, 2, 3, 4, 5, 6, 7], 0) >>> batchify_fn = NamedTuple(SampleData, {'data': Pad(pad_val=0), 'label': Stack()}) >>> sample = batchify_fn([a, b, c]) >>> sample SampleData(data= [[1. 2. 3. 4. 0. 0. 0.] [5. 7. 0. 0. 0. 0. 0.] [1. 2. 3. 4. 5. 6. 7.]] <NDArray 3x7 @cpu_shared(0)>, label= [0 1 0] <NDArray 3 @cpu_shared(0)>) >>> sample.data <BLANKLINE> [[1. 2. 3. 4. 0. 0. 0.] [5. 7. 0. 0. 0. 0. 0.] [1. 2. 3. 4. 5. 6. 7.]] <NDArray 3x7 @cpu_shared(0)> >>> # Let's consider to use a list >>> batchify_fn = NamedTuple(SampleData, [Pad(pad_val=0), Stack()]) >>> batchify_fn([a, b, c]) SampleData(data= [[1. 2. 3. 4. 0. 0. 0.] [5. 7. 0. 0. 0. 0. 0.] [1. 2. 3. 4. 5. 6. 7.]] <NDArray 3x7 @cpu_shared(0)>, label= [0 1 0] <NDArray 3 @cpu_shared(0)>) """ def __init__(self, container, fn_info: t_Union[t_List[t_Callable], t_Tuple[t_Callable], t_Dict[AnyStr, t_Callable]]): self._container = container if isinstance(fn_info, (list, tuple)): if len(container._fields) != len(fn_info): raise ValueError('Attributes mismatch! Required fields={}, fn_info={}' .format(container._fields, fn_info)) elif isinstance(fn_info, dict): for name in container._fields: if name not in fn_info: raise ValueError('Attribute {} has not been assigned a callable. ' 'Required fields={}, Found fields={}' .format(name, container._fields, fn_info.keys())) if len(container._fields) != len(fn_info): raise ValueError('Attributes mimatch! Required fields={}, Found fields={}' .format(container._fields, fn_info.keys())) fn_info = [fn_info[name] for name in container._fields] for fn in fn_info: if not hasattr(fn, '__call__'): raise ValueError('All batchify functions must be callable.') self._fn_l = fn_info
[docs] def __call__(self, data): """Batchify the input data. Parameters ---------- data : List of NamedTuple The samples to batchify. Each sample should be a NamedTuple. Returns ------- ret : List of NamedTuple A namedtuple of length N. Contains the batchified result of each attribute in the input. """ if not isinstance(data[0], self._container): raise ValueError('The samples should have the same type as the stored namedtuple.' ' data[0]={}, container={}'.format(data[0], self._container)) ret = [] for i, ele_fn in enumerate(self._fn_l): ret.append(ele_fn([ele[i] for ele in data])) return self._container(*ret)