# -*- coding: utf-8 -*- # This file is part of MegBrain. # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. """This python module contains functions to apply the operators defined by megbrain. .. note:: Most of the functions are automatically generated, and their signature have the form contain a ``param`` argument (or more than one arguments such as :func:`convolution` that has ``param`` and ``execution_polity``) and also accept keyword arguments. In such case, it can be called by either providing a param object of appropriate type, or by passing the arguments needed by the constructor of param object to the keyword arguments. Furthermore, for a param that needs an enumeration member, the enum name can be used to refer to the enum object. For example, the following statements are equivalent:: elemwise([a, b], mode='max') elemwise([a, b], mode=opr_param_defs.Elemwise.Mode.MAX) elemwise([a, b], param=opr_param_defs.Elemwise('max')) """ from . import mgb as _mgb from . import helper as _helper from . import opr_param_defs as _opr_param_defs import sys import enum import collections import json __git_commit__ = "{%git_commit%}" {%body%} class _ElemMeta(type): def __getattr__(self, name): def run(*inputs, **kwargs): return elemwise(inputs, mode=name, **kwargs) if name.startswith('__'): return return run class elem(metaclass=_ElemMeta): """ Helper class for easily applying element-wise operator. Request for getting member method would be translated to a call to :func:`elemwise` with mode set to the method name. Example:: elem.exp(a) # elemwise(a, mode='exp') elem.max(a, b) # elemwise([a, b], mode='max') """ def add_update( dest, delta, alpha=_mgb.SharedScalar(1), beta=_mgb.SharedScalar(1), bias=_mgb.SharedScalar(0), disable=_mgb.SharedScalar(0), *, name=None, comp_node=None, config=None, comp_graph=None): """update *dest* by `dest := dest*alpha + delta*beta + bias` :param dest: target var to be updated; must be created from :func:`make_shared` :type dest: :class:`.SymbolVar` :param disable: AddUpdate will not be executed if disable is set to 1, this is used for dynamic param-updating. The value of this SharedScalar can only be set to 0/1 of type `int` :type disable: :class:`.SharedScalar` """ def as_ss(x): if not isinstance(x, _mgb.SharedScalar): x = _mgb.SharedScalar(x) return x assert isinstance(dest, _mgb.SymbolVar) config = _helper.gen_config(name, comp_node, config) dest, delta = _helper.canonize_input_vars( [dest, delta], comp_graph=comp_graph, config=config) assert isinstance(disable, _mgb.SharedScalar) alpha, beta, bias = map(as_ss, [alpha, beta, bias]) return _mgb._Opr.add_update(dest, delta, alpha, beta, bias, disable, config) def reduce_(src, mode, axis=None, keepdims=False, *, name=None, comp_node=None, config=None, comp_graph=None): """reduce along given axis; if axis is None, reduce to scalar :param mode: reduction mode :type mode: :class:`~megengine._internal.opr_param_defs.Reduce.Mode` compatible :param axis: axis along which to reduce input var :type axis: int :param keepdims: whether to keep an axis of shape 1 in the result :type keepdims: False """ assert isinstance(src, _mgb.SymbolVar) config = _helper.gen_config(name, comp_node, config) inputs = [src] kwargs = {'mode': mode} remove_axis = False if axis is None: inputs.append(1) assert not keepdims, 'can not set axis=None and keepdims=True' else: remove_axis = not keepdims kwargs['axis'] = axis ret = reduce_general(inputs, config=config, comp_graph=comp_graph, **kwargs) if remove_axis: ret = _mgb._Opr.remove_axis(ret, axis, _mgb.make_opr_config()) return _helper.cvt_opr_result(ret) def _reduce_like(impl, src, axis, keepdims, name, comp_node, config, comp_graph): config = _helper.gen_config(name, comp_node, config) remove_axis = False if axis is None: assert not keepdims, 'can not set axis=None and keepdims=True' src = src.flatten() axis = 0 else: assert isinstance(axis, int) and axis >= 0, ( 'bad axis: {!r}'.format(axis)) remove_axis = not keepdims ret = impl(src, axis=axis, config=config, comp_graph=comp_graph) if remove_axis: ret = _mgb._Opr.remove_axis(ret, axis, _mgb.make_opr_config()) return _helper.cvt_opr_result(ret) def dimshuffle(src, pattern, ndim=0, *, name=None, comp_node=None, config=None): """swap shapes and strides according to given pattern :param pattern: a list of integers, where each element is the input axis of that output axis. An element could also be 'x' for creating a new axis with shape 1 :param ndim: number of input dimensions; 0 to be inferred from pattern; this is required only for grad """ config = _helper.gen_config(name, comp_node, config) if not isinstance(pattern, (list, tuple)): raise TypeError('could not convert {} to dimshuffle pattern'.format( pattern)) pattern_mgb = _mgb._VectorInt() for i in pattern: if i == 'x': pattern_mgb.push_back(-1) else: i = int(i) assert i >= 0 pattern_mgb.push_back(i) return _mgb._Opr.dimshuffle(src, pattern_mgb, int(ndim), config) def param_pack_split(src, shapes, *, name=None, comp_node=None, config=None): """ split param into a list of tensor for given shape ParamPackSplit operator has a input: ``src`` and would have a ``output``. output[i] indicates the address of tensor which part of ``src`` would transfer its elements into. Example: a input tensor with size 32, the shapes: ``[(1, 2, 4), (4, 2, 2), (4, 2, 1)]``, the output tensor would be a list of address with size 3. output[0] indicates the address of tensor with shapes[0]:(1, 2, 4), output[1] indicates the address of tensor with shapes[1]:(4, 2, 2), output[2] indicates the address of tensor with shapes[2]:(4, 2, 1). :param src: The concatenated input tensor. :type src: :class:`SymbolVar` :param shapes: Shapes of output tensors :type shapes: list of list of int """ config = _helper.gen_config(name, comp_node, config) if not isinstance(shapes, (list, tuple)): raise TypeError('could not convert {} to tensor shapes'.format( shapes)) shapes_mgb = _mgb._VectorTensorShape() for s in shapes: s = tuple(map(int, s)) assert min(s) > 0 shapes_mgb.push_back(s) return _mgb._Opr.param_pack_split(src, shapes_mgb, config) class _modify_subtensor_helper: def __init__(self, dest, val, *, name=None, comp_node=None, config=None): self.dest = dest self.val = val self.config = _helper.gen_config(name, comp_node, config) def __getitem__(self, idx): inp = _mgb._VectorSymbolVar() dest, desc = _helper.cvt_getitem_to_idx_desc( self.dest, idx, allow_newaxis=False) assert desc is not None, 'no __getitem__ entries given' inp.push_back(dest) inp.push_back(self.val) return _mgb._create_subtensor_like_opr( self._opr_name, inp, desc, self.config) class set_subtensor(_modify_subtensor_helper): """a proxy object which supports ``__getitem__`` to set subtensor. ``c = set_subtensor(a, b)[idx]`` is equivalent to the numpy expression:: c = a.copy() c[idx] = b """ _opr_name = 'set_subtensor' class incr_subtensor(_modify_subtensor_helper): """a proxy object which supports ``__getitem__`` to increase subtensor. ``c = incr_subtensor(a, b)[idx]`` is equivalent to the numpy expression:: c = a.copy() c[idx] += b """ _opr_name = 'incr_subtensor' class mesh_indexing: """ Extract elements from given tensor by the coordinates which is Cartesian product of given index; example:: mesh_indexing(x)[:, [2, 3], :, [2, 3, 4]] """ def __init__(self, src, *, name=None, comp_node=None, config=None): self.src = src self.config = _helper.gen_config(name, comp_node, config) def __getitem__(self, idx): inp, desc = _helper.cvt_getitem_to_idx_desc(self.src, idx) if desc is None: return inp return _mgb._create_subtensor_like_opr( 'mesh_indexing', [inp], desc, self.config) class batched_mesh_indexing: """ Similar to :class:`mesh_indexing`, while the k-th position of slices is a 2-dim matrix `matrix[k]`. The `matrix[k] is a list of index. The i-th row `matrix[k][i]` represents the index of the associated k-th position slice when `batch_idx == i` ; example:: batched_mesh_indexing(x)[:, [[1, 2], [2, 3]], 1:-1:-1] .. warning:: The first dimension of slices must be (start, stop, step) like, cannot be any of SymbolVar, numpy.array, Python list. And the shape of other indexs must be (n, x) while n is the length of first dimension of tensor after applying [start:stop:step] """ def __init__(self, src, *, name=None, comp_node=None, config=None): self.src = src self.config = _helper.gen_config(name, comp_node, config) def __getitem__(self, idx): inp, desc = _helper.cvt_getitem_to_idx_desc(self.src, idx) if desc is None: return inp return _mgb._create_subtensor_like_opr( 'batched_mesh_indexing', [inp], desc, self.config) class incr_mesh_indexing(_modify_subtensor_helper): _opr_name = 'incr_mesh_indexing' class set_mesh_indexing(_modify_subtensor_helper): _opr_name = 'set_mesh_indexing' class batched_incr_mesh_indexing(_modify_subtensor_helper): _opr_name = 'batched_incr_mesh_indexing' class batched_set_mesh_indexing(_modify_subtensor_helper): _opr_name = 'batched_set_mesh_indexing' class advanced_indexing: """wrapper for numpy-like advanced indexing, where a non-slice index can be a vector; example:: advanced_indexing(x)[:, [2, 3]] """ def __init__(self, src, *, name=None, comp_node=None, config=None): self.src = src self.config = _helper.gen_config(name, comp_node, config) def __getitem__(self, idx): inp, desc = _helper.cvt_getitem_to_idx_desc(self.src, idx) if desc is None: return inp return _mgb._create_subtensor_like_opr( 'mavi', [inp], desc, self.config) class set_advanced_indexing(_modify_subtensor_helper): """:class:`set_subtensor` equivalent with advanced-indexing support""" _opr_name = 'set_mavi' class incr_advanced_indexing(_modify_subtensor_helper): """:class:`incr_subtensor` equivalent with advanced-indexing support""" _opr_name = 'incr_mavi' def mean(inp, axis, keepdims): """average value along an axis""" if hasattr(inp.dtype, 'metadata'): return reduce_(inp, 'MEAN', axis, keepdims) else: s = reduce_(inp, 'SUM', axis, keepdims) if axis is None: cnt = inp.shape.prod() else: cnt = inp.axis_shape(axis) return s / cnt def square(inp): """*inp* squared""" return inp ** 2 def sqrt(inp): """square root""" return inp ** 0.5 class _LoopDescMakerCallback(_mgb._LoopDescMakerCallback): def __init__(self, func): super().__init__() assert isinstance(func, collections.Callable) self._func = func self.__disown__() def call(self, desc): self._func(desc) def make_loop(desc_maker, *, swap_interval=-5, name=None, comp_node=None, config=None): """Create a loop operator. The loop operator works in the following way: 1. Copy variables specified by :meth:`.LoopDesc.add_input` from the parent graph into the sub graph. 2. Evaluates the loop condition. 3. If the absolute value of the loop condition is no more than 1e-6, go to 5. 4. Update variables in the sub graph using rules specified by :meth:`.LoopDesc.assign` and then go to 2 again. 5. Copy values of output variables given by :meth:`.LoopDesc.add_output` into the parent graph and exit. The loop operator could be thought of as a digital circuit, where the sub graph (which must be purely functional) is the combinational logic part and the :meth:`.LoopDesc.assign` rules serve as the flip-flops. :type desc_maker: callable :param desc_maker: a function to create the loop descriptor; it would receive a :class:`.LoopDesc` object and should call methods on it to describe the sub graph. This function may be called multiple times, and it should behave exactly the same in every call. :type swap_interval: int :param swap_interval: number of loop executions between swapping saved mutable states to host; larger *swap_interval* requires more memory and less copy stall. If *swap_interval* is negative, then statically inferred loop time would be used if possible; otherwise its absolute value would be used as swap interval. :rtype: list of :class:`.SymbolVar` :return: the output vars, corresponding to each :meth:`.LoopDesc.add_output` call. """ config = _helper.gen_config(name, comp_node, config) return _mgb._make_loop(_LoopDescMakerCallback(desc_maker), swap_interval, config) def symvar_from_shared_nd(sv, comp_graph, name=None): """get a symbol var in a computing graph that represents a shared (i.e. pre-allocated) value on device :param sv: the shared value :type sv: :class:`.SharedND` :param comp_graph: the computing graph to which this symvar should belong :type graph: :class:`.CompGraph` :param name: the name of resulting symvar :type name: str or None :rtype: :class:`.SymbolVar` """ assert isinstance(sv, _mgb.SharedND) return sv.symvar(comp_graph, name) def zero_grad(sv, **kwargs): return set_grad(sv, None, **kwargs) # for backward pickle compatiblility def _make_enum_unpickle(new_enum): """create a class that can be used for unpickling old enum values""" class OldEnum: def __new__(cls, value): return new_enum[value] return OldEnum ConvMode = _make_enum_unpickle(_opr_param_defs.Convolution.Mode) PoolingMode = _make_enum_unpickle(_opr_param_defs.Pooling.Mode) ROIPoolingMode = _make_enum_unpickle(_opr_param_defs.ROIPooling.Mode) WarpPerspectiveBorderMode = _make_enum_unpickle( _opr_param_defs.WarpPerspective.BorderMode) WarpPerspectiveInterpMode = _make_enum_unpickle( _opr_param_defs.WarpPerspective.InterpolationMode)