Python mxnet.Context() Examples

def __init__(
        self,
        dataset: Dataset,
        *,
        transform: Transformation,
        batch_size: int,
        ctx: mx.Context,
        num_workers: Optional[int] = None,
        num_prefetch: Optional[int] = None,
        dtype: DType = np.float32,
        **kwargs,
    ) -> None:
        super().__init__(
            dataset=dataset,
            transform=transform,
            is_train=True,
            batch_size=batch_size,
            ctx=ctx,
            dtype=dtype,
            cyclic=False,
            num_workers=num_workers,
            num_prefetch=num_prefetch,
            shuffle_buffer_length=None,
            **kwargs,
        ) 

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

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

def train(train_iter, test_iter, net, loss, trainer, ctx, num_epochs):
    """Train and evaluate a model."""
    print('training on', ctx)
    if isinstance(ctx, mx.Context):
        ctx = [ctx]
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n, m, start = 0.0, 0.0, 0, 0, time.time()
        for i, batch in enumerate(train_iter):
            Xs, ys, batch_size = _get_batch(batch, ctx)   
            with autograd.record():
                y_hats = [net(X) for X in Xs]
                ls = [loss(y_hat, y) for y_hat, y in zip(y_hats, ys)]
            for l in ls:
                l.backward()
            trainer.step(batch_size)
            train_l_sum += sum([l.sum().asscalar() for l in ls])
            n += sum([l.size for l in ls])
            train_acc_sum += sum([(y_hat.argmax(axis=1) == y).sum().asscalar()
                                 for y_hat, y in zip(y_hats, ys)])
            m += sum([y.size for y in ys])
        test_acc = evaluate_accuracy(test_iter, net, ctx)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f, '
              'time %.1f sec'
              % (epoch + 1, train_l_sum / n, train_acc_sum / m, test_acc,
                 time.time() - start)) 

def _stack_arrs(arrs, use_shared_mem=False):
        """
        Internal imple for stacking arrays.
        """
        if isinstance(arrs[0], mx.nd.NDArray):
            if use_shared_mem:
                out = mx.nd.empty((len(arrs),) + arrs[0].shape, dtype=arrs[0].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)
            if use_shared_mem:
                return mx.nd.array(out, ctx=mx.Context("cpu_shared", 0))
            else:
                return mx.nd.array(out) 

def train(train_iter, test_iter, net, loss, trainer, ctx, num_epochs):
    """Train and evaluate a model."""
    print('training on', ctx)
    if isinstance(ctx, mx.Context):
        ctx = [ctx]
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n, m, start = 0.0, 0.0, 0, 0, time.time()
        for i, batch in enumerate(train_iter):
            Xs, ys, batch_size = _get_batch(batch, ctx)   
            with autograd.record():
                y_hats = [net(X) for X in Xs]
                ls = [loss(y_hat, y) for y_hat, y in zip(y_hats, ys)]
            for l in ls:
                l.backward()
            trainer.step(batch_size)
            train_l_sum += sum([l.sum().asscalar() for l in ls])
            n += sum([l.size for l in ls])
            train_acc_sum += sum([(y_hat.argmax(axis=1) == y).sum().asscalar()
                                 for y_hat, y in zip(y_hats, ys)])
            m += sum([y.size for y in ys])
        test_acc = evaluate_accuracy(test_iter, net, ctx)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f, '
              'time %.1f sec'
              % (epoch + 1, train_l_sum / n, train_acc_sum / m, test_acc,
                 time.time() - start)) 

def _stack_arrs(arrs, use_shared_mem=False):
        """
        Internal imple for stacking arrays.
        """
        if isinstance(arrs[0], mx.nd.NDArray):
            if use_shared_mem:
                out = mx.nd.empty((len(arrs),) + arrs[0].shape, dtype=arrs[0].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)
            if use_shared_mem:
                return mx.nd.array(out, ctx=mx.Context("cpu_shared", 0))
            else:
                return mx.nd.array(out) 

def benchmarking(net, opt, ctx):
    if isinstance(ctx, mx.Context):
        ctx = [ctx]

    bs = opt.batch_size
    num_iterations = opt.num_iterations
    input_shape = (bs, 3,) + tuple(input_size)
    size = num_iterations * bs
    data = mx.random.uniform(-1.0, 1.0, shape=input_shape, ctx=ctx[0], dtype='float32')
    dry_run = 5

    from tqdm import tqdm
    with tqdm(total=size + dry_run * bs) as pbar:
        for n in range(dry_run + num_iterations):
            if n == dry_run:
                tic = time.time()
            output = net(data)
            output.wait_to_read()
            pbar.update(bs)
    speed = size / (time.time() - tic)
    print('With batch size %d , %d batches, throughput is %f imgs/sec' % (bs, num_iterations, speed)) 

def batchify(
    data: List[dict],
    dtype: DType,
    multi_processing: bool,
    single_process_ctx: Optional[mx.Context] = None,
    variable_length: bool = False,
) -> DataBatch:
    """reduce the list of dictionaries to a single dictionary, where values
        referenced by identical key are reduced using the stack function"""
    return {
        key: stack(
            data=[item[key] for item in data],
            multi_processing=multi_processing,
            dtype=dtype,
            single_process_ctx=single_process_ctx,
            variable_length=variable_length,
        )
        for key in data[0].keys()
    } 

def _as_in_context(batch: dict, ctx: mx.Context) -> DataBatch:
    """Move data into new context, should only be in main process."""
    assert (
        not MPWorkerInfo.worker_process
    ), "This function is not meant to be used in workers."
    batch = {
        k: v.as_in_context(ctx) if isinstance(v, nd.NDArray)
        # Workaround due to MXNet not being able to handle NDArrays with 0 in shape properly:
        else (
            stack(v, False, v.dtype, ctx)
            if isinstance(v[0], np.ndarray) and 0 in v[0].shape
            else v
        )
        for k, v in batch.items()
    }
    return batch 

def __init__(
        self,
        dataset: Dataset,
        *,
        transform: Transformation,
        batch_size: int,
        ctx: mx.Context,
        num_workers: Optional[int] = None,
        num_prefetch: Optional[int] = None,
        dtype: DType = np.float32,
        **kwargs,
    ) -> None:
        super().__init__(
            dataset=dataset,
            transform=transform,
            is_train=False,
            batch_size=batch_size,
            ctx=ctx,
            dtype=dtype,
            cyclic=False,
            num_workers=num_workers,
            num_prefetch=num_prefetch,
            shuffle_buffer_length=None,
            **kwargs,
        ) 

def get_gpu_memory_usage(ctx: Union[mx.context.Context, List[mx.context.Context]]) -> Dict[int, Tuple[int, int]]:
    """
    Returns used and total memory for GPUs identified by the given context list.

    :param ctx: List of MXNet context devices.
    :return: Dictionary of device id mapping to a tuple of (memory used, memory total).
    """
    if isinstance(ctx, mx.context.Context):
        ctx = [ctx]
    ctx = [c for c in ctx if c.device_type == 'gpu']
    if not ctx:
        return {}

    memory_data = {}  # type: Dict[int, Tuple[int, int]]
    for c in ctx:
        try:
            free, total = mx.context.gpu_memory_info(device_id=c.device_id)  # in bytes
            used = total - free
            memory_data[c.device_id] = (used * 1e-06, total * 1e-06)
        except mx.MXNetError:
            logger.exception("Failed retrieving memory data for gpu%d", c.device_id)
            continue
    log_gpu_memory_usage(memory_data)
    return memory_data 

def seed_rngs(seed: int, ctx: Optional[Union[mx.Context, List[mx.Context]]] = None) -> None:
    """
    Seed the random number generators (Python, Numpy and MXNet).

    :param seed: The random seed.
    :param ctx: Random number generators in MXNet are device specific.
           If None, MXNet will set the state of each generator of each device using seed and device id. This will lead
           to different results on different devices. If ctx is provided, this function will seed
           device-specific generators with a fixed offset. E.g. for 2 devices and seed=13, seed for gpu(0) will be 13,
           14 for gpu(1). See https://beta.mxnet.io/api/gluon-related/_autogen/mxnet.random.seed.html.
    """
    logger.info("Random seed: %d", seed)
    np.random.seed(seed)
    random.seed(seed)
    if ctx is None:
        mx.random.seed(seed, ctx='all')
    else:
        if isinstance(ctx, mx.Context):
            ctx = [ctx]
        for i, c in enumerate(ctx):
            mx.random.seed(seed + i, ctx=c) 

def initialize_from_array(
        self, input_array: np.ndarray, ctx: mx.Context = get_mxnet_context()
    ):
        r"""
        Initialize the representation based on a numpy array.

        Parameters
        ----------
        input_array
            Numpy array.
        ctx
            MXNet context.
        """
        pass

    # noinspection PyMethodOverriding 

def initialize_from_dataset(
        self, input_dataset: Dataset, ctx: mx.Context = get_mxnet_context()
    ):
        self.initialize_from_array(np.array([]), ctx) 

def initialize_from_array(
        self, input_array: np.ndarray, ctx: mx.Context = get_mxnet_context()
    ):
        for representation in self.chain:
            representation.initialize_from_array(input_array, ctx)

    # noinspection PyMethodOverriding 

def get_vgg_atrous_extractor(num_layers, im_size, pretrained=False, ctx=mx.cpu(),
                             root=os.path.join('~', '.mxnet', 'models'), **kwargs):
    """Get VGG atrous feature extractor networks.

    Parameters
    ----------
    num_layers : int
        VGG types, can be 11,13,16,19.
    im_size : int
        VGG detection input size, can be 300, 512.
    pretrained : bool
        Load pretrained weights if `True`.
    ctx : mx.Context
        Context such as mx.cpu(), mx.gpu(0).
    root : str
        Model weights storing path.

    Returns
    -------
    mxnet.gluon.HybridBlock
        The returned network.

    """
    layers, filters = vgg_spec[num_layers]
    extras = extra_spec[im_size]
    net = VGGAtrousExtractor(layers, filters, extras, **kwargs)
    if pretrained:
        from ..model_store import get_model_file
        batch_norm_suffix = '_bn' if kwargs.get('batch_norm') else ''
        net.initialize(ctx=ctx)
        net.load_params(get_model_file('vgg%d_atrous%s'%(num_layers, batch_norm_suffix),
                                       root=root), ctx=ctx, allow_missing=True)
    return net 

def initialize_from_dataset(
        self, input_dataset: Dataset, ctx: mx.Context = get_mxnet_context()
    ):
        r"""
        Initialize the representation based on an entire dataset.

        Parameters
        ----------
        input_dataset
            GluonTS dataset.
        ctx
            MXNet context.
        """
        pass 

def initialize_from_array(
        self, input_array: np.ndarray, ctx: mx.Context = get_mxnet_context()
    ):
        for representation in self.representations:
            representation.initialize_from_array(input_array, ctx)

    # noinspection PyMethodOverriding 

def initialize_from_dataset(
        self, input_dataset: Dataset, ctx: mx.Context = get_mxnet_context()
    ):
        for representation in self.chain:
            representation.initialize_from_dataset(input_dataset, ctx) 

def encode_mx_context(v: mx.Context) -> Any:
    """
    Specializes :func:`encode` for invocations where ``v`` is an instance of
    the :class:`~mxnet.Context` class.
    """
    return {
        "__kind__": kind_inst,
        "class": fqname_for(v.__class__),
        "args": encode([v.device_type, v.device_id]),
    } 

def _stack_arrs(arrs, use_shared_mem=False):
    """Internal imple for stacking arrays."""
    if isinstance(arrs[0], mx.nd.NDArray):
        if use_shared_mem:
            out = mx.nd.empty((len(arrs),) + arrs[0].shape, dtype=arrs[0].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)
        if use_shared_mem:
            return mx.nd.array(out, ctx=mx.Context('cpu_shared', 0))
        else:
            return mx.nd.array(out) 

def __init__(self, data_shape, label_shape, labels_range, dtype):
        """MXNet partitions data batch evenly among the available GPUs. Here, the
           batch size is the effective batch size.
           
        Memory for data and label tensors is allocated with `cpu_pinned` context.
        To make this iterator consistent with other iterators that provide real data,
        I always set maximal number of iterations to be `warmup_iters + bench_iters`.

        Changes:
            New version does not support `max_iter` parameter. To limit number of batches, wrap this iterator with
                `mx.io.ResizeIter`.

        Args:
            data_shape (tuple): Shape of input data tensor (X) including batch size. The batch size is the 0th
                dimension (bsz = data_shape[0]). This batch size must be an effective batch for a whole node.
            label_shape (tuple): Shape of input label tensor (Y) including batch size. The batch size is the 0th
                dimension (bsz = labels_shape[0]). This batch size must be an effective batch for a whole node.
            labels_range (list): List of output labels. For ImageNet, that would be a list with integers from 0 to 999.
            dtype (str): Data type for data tensor (float32, float16).
        """
        super(SyntheticDataIterator, self).__init__(data_shape[0])
        # Let's assume this data iterator always returns single precision tensors.
        self.dtype = dtype
        # mx.Context: cpu, gpu, cpu_pinned, cpu_shared
        # It seems in latest NGC containers this bug is present: https://github.com/apache/incubator-mxnet/pull/12031
        # The way how mxnet sends data to GPUs, it's better to keep all in CPU. Other possible option would be to
        # eliminate CPU memory at all with synthetic data which is not probably a great idea.
        synth_data_context = os.environ.get('DLBS_MXNET_SYNTHETIC_DATA_CONTEXT', 'cpu')
        self.data = mx.nd.array(np.random.uniform(-1, 1, data_shape),
                                dtype=self.dtype,
                                ctx=mx.Context(synth_data_context, 0))
        self.label_shape = [label_shape[0]]
        if not self.label_shape:
            self.label_shape = [self.batch_size]
        self.label = mx.nd.array(np.random.randint(labels_range[0], labels_range[1] + 1, self.label_shape),
                                 dtype='float32',
                                 ctx=mx.Context(synth_data_context, 0)) 

def __init__(self, num_classes, data_shape, max_iter, dtype):
        self.batch_size = data_shape[0]
        self.cur_iter = 0
        self.max_iter = max_iter
        self.dtype = dtype
        label = np.random.randint(0, num_classes, [self.batch_size,])
        data = np.random.uniform(-1, 1, data_shape)
        self.data = mx.nd.array(data, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0))
        self.label = mx.nd.array(label, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0)) 

def load_parameters(self,
                        filename: str,
                        ctx: Union[mx.Context, List[mx.Context]] = None,
                        allow_missing: bool = False,
                        ignore_extra: bool = False,
                        cast_dtype: bool = False,
                        dtype_source: str = 'current'):
        """Load parameters from file previously saved by `save_parameters`.

        Parameters
        ----------
        filename : str
            Path to parameter file.
        ctx : Context or list of Context, default cpu()
            Context(s) to initialize loaded parameters on.
        allow_missing : bool, default False
            Whether to silently skip loading parameters not represents in the file.
        ignore_extra : bool, default False
            Whether to silently ignore parameters from the file that are not
            present in this Block.
        cast_dtype : bool, default False
            Cast the data type of the NDArray loaded from the checkpoint to the dtype
            provided by the Parameter if any.
        dtype_source : str, default 'current'
            must be in {'current', 'saved'}
            Only valid if cast_dtype=True, specify the source of the dtype for casting
            the parameters
        References
        ----------
        `Saving and Loading Gluon Models \
        <https://mxnet.incubator.apache.org/tutorials/gluon/save_load_params.html>`_
        """
        utils.check_condition(os.path.exists(filename), "No model parameter file found under %s. "
                                                     "This is either not a model directory or the first training "
                                                     "checkpoint has not happened yet." % filename)
        super().load_parameters(filename, ctx=ctx, allow_missing=allow_missing, ignore_extra=ignore_extra,
                                cast_dtype=cast_dtype, dtype_source=dtype_source)
        logger.info('Loaded params from "%s" to "%s"', filename, mx.cpu() if ctx is None else ctx) 

def determine_context(device_ids: List[int],
                      use_cpu: bool,
                      disable_device_locking: bool,
                      lock_dir: str,
                      exit_stack: ExitStack) -> List[mx.Context]:
    """
    Determine the MXNet context to run on (CPU or GPU).

    :param device_ids: List of device as defined from the CLI.
    :param use_cpu: Whether to use the CPU instead of GPU(s).
    :param disable_device_locking: Disable Sockeye's device locking feature.
    :param lock_dir: Directory to place device lock files in.
    :param exit_stack: An ExitStack from contextlib.

    :return: A list with the context(s) to run on.
    """
    if use_cpu:
        context = [mx.cpu()]
    else:
        num_gpus = get_num_gpus()
        check_condition(num_gpus >= 1,
                        "No GPUs found, consider running on the CPU with --use-cpu ")
        if horovod_mpi.using_horovod():
            # Running with Horovod/MPI: GPU(s) are determined by local rank
            check_condition(len(device_ids) == 1 and device_ids[0] < 0,
                            "When using Horovod, --device-ids should be a negative integer indicating the number of "
                            "GPUs each worker should use.")
            n_ids = -device_ids[0]
            context = [mx.gpu(_id + horovod_mpi.hvd.local_rank() * n_ids) for _id in range(n_ids)]
        else:
            if disable_device_locking:
                context = expand_requested_device_ids(device_ids)
            else:
                context = exit_stack.enter_context(acquire_gpus(device_ids, lock_dir=lock_dir))
            context = [mx.gpu(gpu_id) for gpu_id in context]
    return context 

def __init__(self, num_classes, data_shape, max_iter, dtype):
        self.batch_size = data_shape[0]
        self.cur_iter = 0
        self.max_iter = max_iter
        self.dtype = dtype
        label = np.random.randint(0, num_classes, [self.batch_size,])
        data = np.random.uniform(-1, 1, data_shape)
        self.data = mx.nd.array(data, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0))
        self.label = mx.nd.array(label, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0)) 

def __init__(self, num_classes, data_shape, max_iter, dtype):
        self.batch_size = data_shape[0]
        self.cur_iter = 0
        self.max_iter = max_iter
        self.dtype = dtype
        label = np.random.randint(0, num_classes, [self.batch_size,])
        data = np.random.uniform(-1, 1, data_shape)
        self.data = mx.nd.array(data, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0))
        self.label = mx.nd.array(label, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0)) 

def __init__(self, ctx):
        super(CrossEntropyLoss, self).__init__()
        self.eps = 1e-6 # Avoid -inf when taking log(0)
        self.eps1 = 1. + self.eps
        self.eps_1 = 1. - self.eps
        self.ctx=mx.Context(ctx) 

def __init__(self, num_classes, data_shape, max_iter, dtype):
        self.batch_size = data_shape[0]
        self.cur_iter = 0
        self.max_iter = max_iter
        self.dtype = dtype
        label = np.random.randint(0, num_classes, [self.batch_size,])
        data = np.random.uniform(-1, 1, data_shape)
        self.data = mx.nd.array(data, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0))
        self.label = mx.nd.array(label, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0)) 

def test_ndarray_copy():
    c = mx.nd.array(np.random.uniform(-10, 10, (10, 10)))
    d = c.copyto(mx.Context('cpu', 0))
    assert np.sum(np.abs(c.asnumpy() != d.asnumpy())) == 0.0