python/mxnet/gluon/data/batchify.py - master - githubmirror/incubator-mxnet

python/mxnet/gluon/data/batchify.py
415 строк14 KB
hankaj
Python string formatting (#21136)
16 сен 2022, 14:01
16 сен 2022, 14:01c8922fe
# 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.

# coding: utf-8
# pylint: disable=reimported, consider-using-enumerate
"""Batchify function."""
import math
import warnings
import numpy as np

from ...device import Device, cpu
from ... import ndarray as nd
from ... import numpy as _np
from ...util import is_np_array

class Stack(object):
    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.
    Examples
    --------
    >>> from mxnet.gluon.data import batchify
    >>> # Stack multiple lists
    >>> a = [1, 2, 3, 4]
    >>> b = [4, 5, 6, 8]
    >>> c = [8, 9, 1, 2]
    >>> batchify.Stack()([a, b, c])
    [[1. 2. 3. 4.]
     [4. 5. 6. 8.]
     [8. 9. 1. 2.]]
    <NDArray 3x4 @cpu(0)>
    >>> # Stack multiple numpy.ndarrays
    >>> import numpy as np
    >>> a = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])
    >>> b = np.array([[5, 6, 7, 8], [1, 2, 3, 4]])
    >>> batchify.Stack()([a, b])
    [[[1. 2. 3. 4.]
      [5. 6. 7. 8.]]
     [[5. 6. 7. 8.]
      [1. 2. 3. 4.]]]
    <NDArray 2x2x4 @cpu(0)>
    >>> # Stack multiple NDArrays
    >>> import mxnet as mx
    >>> a = nd.array([[1, 2, 3, 4], [5, 6, 7, 8]])
    >>> b = nd.array([[5, 6, 7, 8], [1, 2, 3, 4]])
    >>> batchify.Stack()([a, b])
    [[[1. 2. 3. 4.]
      [5. 6. 7. 8.]]
     [[5. 6. 7. 8.]
      [1. 2. 3. 4.]]]
    <NDArray 2x2x4 @cpu(0)>
    """
    def __init__(self, use_shared_mem=False):
        self._use_shared_mem = use_shared_mem

    def __call__(self, data):
        """Batchify the input data
        Parameters
        ----------
        data : list
            The input data samples
        Returns
        -------
        batch_data : NDArray
        """
        _arr = _np if is_np_array() else nd
        _arr_cls = _arr.ndarray if is_np_array() else _arr.NDArray
        if isinstance(data[0], _arr_cls):
            dtype = data[0].dtype
            if self._use_shared_mem:
                out = _arr.empty((len(data),) + data[0].shape, dtype=dtype,
                                 ctx=Device('cpu_shared', 0))
                return _arr.stack(data, out=out) if is_np_array() else _arr.stack(*data, out=out)
            else:
                return _arr.stack(data) if is_np_array() else _arr.stack(*data)
        elif isinstance(data[0], (tuple, list)):
            data = zip(*data)
            return [self.__call__(i) for i in data]
        else:
            out = np.asarray(data)
            dtype = out.dtype
            if self._use_shared_mem:
                return _arr.array(out, ctx=Device('cpu_shared', 0), dtype=dtype)
            else:
                return _arr.array(out, dtype=dtype)

    def __mx_handle__(self):
        from ._internal import StackBatchify
        return StackBatchify()

def _pad_arrs_to_max_length(arrs, pad_val, use_shared_mem, dtype, round_to=None):
    """Inner Implementation of the Pad batchify
    Parameters
    ----------
    arrs : list
    pad_val : number
    use_shared_mem : bool, default False
    round_to : int

    Returns
    -------
    ret : NDArray
    """
    _arr = _np if is_np_array() else nd
    _arr_cls = _np.ndarray if is_np_array() else nd.NDArray
    if isinstance(arrs[0], _arr_cls):
        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]
        dtype = arrs[0][0].dtype if dtype is None else dtype
    else:
        dtype = arrs[0].dtype if dtype is None else dtype

    ret_shape = list(arrs[0].shape)
    for pad_axis in range(len(ret_shape)):
        curr_lengths = [ele.shape[pad_axis] for ele in arrs]
        max_size = max(curr_lengths)
        if round_to is not None:
            max_size = round_to * math.ceil(max_size / round_to)
        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 == ret_shape[1:]:
            ret[i] = arr
        else:
            slices = [slice(None) for _ in range(arr.ndim)]
            for pad_axis in range(arr.ndim):
                slices[pad_axis] = slice(0, arr.shape[pad_axis])
                assert slices[pad_axis].start != slices[pad_axis].stop
            slices = [slice(i, i + 1)] + slices
            ret[tuple(slices)] = arr


    device = Device('cpu_shared', 0) if use_shared_mem else cpu()
    ret = _arr.array(ret, ctx=device, dtype=dtype)

    return ret


class Pad(object):
    """Pad the input ndarrays along the specific padding axis and stack them to get the output.
    Input of the function will be N samples. Each sample should contain a single element that
    can be 1) numpy.ndarray, 2) mxnet.nd.NDArray, 3) list of numbers.
    You can set the `pad_val` to determine the padding value.

    The arrays will be padded to the largest dimensions(at most 5 dimensions to pad) and then
    stacked to form the final output.

    Parameters
    ----------
    val : float or int, default None
        The padding value.
    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
    --------
    >>> from mxnet.gluon.data import batchify
    >>> # Inputs are multiple lists
    >>> a = [1, 2, 3, 4]
    >>> b = [4, 5, 6]
    >>> c = [8, 2]
    >>> batchify.Pad()([a, b, c])
    [[ 1  2  3  4]
     [ 4  5  6  0]
     [ 8  2  0  0]]
    <NDArray 3x4 @cpu(0)>
    >>> # Also output the lengths
    >>> a = [1, 2, 3, 4]
    >>> b = [4, 5, 6]
    >>> c = [8, 2]
    >>> # Inputs are multiple ndarrays
    >>> import numpy as np
    >>> a = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])
    >>> b = np.array([[5, 8], [1, 2]])
    >>> batchify.Pad(val=-1)([a, b])
    [[[ 1  2  3  4]
      [ 5  6  7  8]]
     [[ 5  8 -1 -1]
      [ 1  2 -1 -1]]]
    <NDArray 2x2x4 @cpu(0)>
    >>> # Inputs are multiple NDArrays
    >>> import mxnet as mx
    >>> a = nd.array([[1, 2, 3, 4], [5, 6, 7, 8]])
    >>> b = nd.array([[5, 8], [1, 2]])
    >>> batchify.Pad(val=-1)([a, b])
    [[[ 1.  2.  3.  4.]
      [ 5.  6.  7.  8.]]
     [[ 5.  8. -1. -1.]
      [ 1.  2. -1. -1.]]]
    <NDArray 2x2x4 @cpu(0)>
    """
    def __init__(self, val=None, dtype=None, round_to=None, use_shared_mem=False):
        self._pad_val = 0 if val is None else val
        self._dtype = dtype
        self._warned = False
        self._round_to = round_to
        self._use_shared_mem = use_shared_mem

    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.

        Parameters
        ----------
        data : List[np.ndarray] or List[List[dtype]] or List[nd.NDArray]
            List of samples to pad and stack.
        Returns
        -------
        batch_data: NDArray
            Data in the minibatch. Shape is (N, ...)
        """
        _arr = _np if is_np_array() else nd
        _arr_cls = _arr.ndarray if is_np_array() else _arr.NDArray
        if isinstance(data[0], _arr_cls) 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], (_arr_cls, np.ndarray, list)):
            padded_arr = _pad_arrs_to_max_length(data, self._pad_val,
                                                 self._use_shared_mem,
                                                 self._dtype, self._round_to)
            return padded_arr
        else:
            raise NotImplementedError(
                "Pad() does not support multiple items, use Group(Pad(), Pad(), ...) instead")

    def __mx_handle__(self):
        from ._internal import PadBatchify
        return PadBatchify(pad_val=self._pad_val, dtype=self._dtype if self._dtype is not None else -1)

def _append_arrs(arrs, use_shared_mem=False, expand=False, batch_axis=0):
    """Internal impl for returning appened arrays as list."""
    _arr = _np if is_np_array() else nd
    if isinstance(arrs[0], _arr.NDArray):
        if use_shared_mem:
            out = [x.as_in_context(Device('cpu_shared', 0)) for x in arrs]
        else:
            out = arrs
    else:
        if use_shared_mem:
            out = [_arr.array(x, ctx=Device('cpu_shared', 0)) for x in arrs]
        else:
            out = [_arr.array(x) for x in arrs]

    # add batch axis
    if expand:
        out = [x.expand_dims(axis=batch_axis) for x in out]
    return out


class Append(object):
    r"""Loosely return list of the input data samples.
    There is no constraint of shape for any of the input samples, however, you will
    only be able to apply single batch operations since the output have different shapes.
    Examples
    --------
    >>> a = [1, 2, 3, 4]
    >>> b = [4, 5, 6]
    >>> c = [8, 2]
    >>> batchify.Append()([a, b, c])
    [
    [[1. 2. 3. 4.]]
    <NDArray 1x4 @cpu_shared(0)>,
    [[4. 5. 6.]]
    <NDArray 1x3 @cpu_shared(0)>,
    [[8. 2.]]
    <NDArray 1x2 @cpu_shared(0)>
    ]
    """

    def __init__(self, expand=True, batch_axis=0, use_shared_mem=False):
        self._expand = expand
        self._batch_axis = batch_axis
        self._use_shared_mem = use_shared_mem

    def __call__(self, data):
        """Batchify the input data.
        Parameters
        ----------
        data : list
            The input data samples
        Returns
        -------
        batch_data : NDArray
        """
        return _append_arrs(data, use_shared_mem=self._use_shared_mem,
                            expand=self._expand, batch_axis=self._batch_axis)

class Group(object):
    """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 `Group` will be applied on the `i`th attribute. For example, each
    data sample is (nd_data, label). You can wrap two batchify functions using
    `Group(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
    --------
    >>> a = ([1, 2, 3, 4], 0)
    >>> b = ([5, 7], 1)
    >>> c = ([1, 2, 3, 4, 5, 6, 7], 0)
    >>> f1, f2 = Group(Pad(val=0),
    ...                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):
        self._handle = None
        if isinstance(fn, (list, tuple)):
            assert len(args) == 0, 'Input pattern not understood. The input of Group can be ' \
                                   'Group(A, B, C) or Group([A, B, C]) or Group((A, B, C)). ' \
                                   f'Received fn={str(fn)}, args={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! ' \
                                                f'type(fn[{i}]) = {str(type(ele_fn))}'

    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)

    def __mx_handle__(self):
        if self._handle  is None:
            from ._internal import GroupBatchify
            try:
                mx_fn = [fn.__mx_handle__() for fn in self._fn]
                self._handle = GroupBatchify(functions=mx_fn)
            except Exception as e:
                raise NotImplementedError(
                    "GroupBatchify requires all internal batchify functions supported by backend."
                    + str(e))
        return self._handle

class AsList(object):
    """Simply forward the list of input data.
    This is particularly useful when the Dataset contains textual data
    and in conjonction with the `Group` batchify function.
    Examples
    --------
    >>> a = ([1, 2, 3, 4], "I am using MXNet")
    >>> b = ([5, 7, 2, 5], "Gluon rocks!")
    >>> c = ([1, 2, 3, 4], "Batchification!")
    >>> _, l = Group(Stack(), AsList())([a, b, c])
    >>> l
    ['I am using MXNet', 'Gluon rocks!', 'Batchification!']
    """
    def __call__(self, data):
        """
        Parameters
        ----------
        data : list
            The list of samples
        Returns
        -------
        ret : list
            The input list
        """
        return list(data)