megengine.tensor 源代码

# -*- coding: utf-8 -*-
from typing import Union

import numpy as np

from .core._imperative_rt import CompNode
from .core._imperative_rt.core2 import FormatType
from .core._imperative_rt.core2 import Tensor as _Tensor
from .core._imperative_rt.core2 import apply, set_py_tensor_type
from .core._trace_option import use_symbolic_shape
from .core._wrap import as_device
from .core.ops.builtin import Borrow, Copy, GetVarShape
from .core.tensor.array_method import ArrayMethodMixin
from .device import _valid_device, get_default_device
from .logger import get_logger
from .utils.deprecation import deprecated

logger = get_logger(__name__)


[文档]class Tensor(_Tensor, ArrayMethodMixin): r"""A tensor object represents a multidimensional, homogeneous array of fixed-size items. Tensor is the primary MegEngine data structure. Data type(dtype) describes the format of each element, such as ``float32``, ``int8`` and so on, see :ref:`tensor-dtype` for more details. It is similar to :class:`numpy.ndarray` but not the same in the design. For example, GPU devices can be used to store Tensors and execute calculations in MegEngine. The concept of `view <https://numpy.org/doc/stable/reference/generated/numpy.ndarray.view.html>`_ does not exist in MegEngine so indexing and other behaviors might be different with NumPy. All manipulations and operations on/between Tensors could be found in the :mod:`~.megengine.functional` module. Keep in mind that they are **not in-place**, a new Tensor will always be returned and the original data will remain constant. For more information, refer to the :ref:`tensor-guide` topic. Args: data(Tensor, :class:`~.numpy.ndarray`, :class:`list` or Python number): The data used for construcing Tensor. Tensor could be constructed from a Python :class:`list` / :class:`tuple` or sequence; a NumPy :class:`~.numpy.ndarray` data structure; MegEngine builtin methods and so on. Refer to :ref:`tensor-creation` for more details. dtype(:attr:`~.Tensor.dtype`): The data type of returned Tensor. Infer from ``data`` if not specified. device(:attr:`~.Tensor.device`): The desired device of returned Tensor. Uses :func:`get_default_device` if not specified. is_const: Whether make it a ``ImutableTensor`` in tracing mode, refer to :class:`.jit.trace`. no_cache: Whether cache it for memory sharing. name: Used to improve convenience in graph operation on dumped model. format: Used to indicate which memory format Tensor uses. It will not affect actual memory order or stride, but may affect some operators related to indexing and dimension. Only support "default", "nchw" and "nhwc". .. note:: There are some methods like :meth:`~.Tensor.reshape` / :meth:`~.Tensor.flatten` / :meth:`~.Tensor.transpose` / :meth:`~.Tensor.min` / :meth:`~.Tensor.max` / :meth:`~.Tensor.mean` / :meth:`~.Tensor.sum` / :meth:`~.Tensor.prod` implemented in ``Tensor`` class for convenience and historical reasons. But other methods implemented in the :mod:`~.megengine.functional` module will not be added here anymore, it is hard for maintaining and too many candidates will affect code completion experience. """ grad = None dmap_callback = None _qparams = None _custom_name = "" _name = None _short_name = None _prefix = None def __init__( self, data: Union["Tensor", np.ndarray, list, int, float], dtype: np.dtype = None, device: str = None, is_const: bool = False, no_cache: bool = False, name: str = None, format: str = "default", ): if name is None: name = "" else: self._set_name(name) self._custom_name = name self._name = name self._short_name = name self._prefix = None @property def shape(self) -> Union[tuple, "Tensor"]: r"""Returns a :class:`tuple` or a :class:`~.Tensor` represents tensor dimensions. Note: The shape of a tensor was usually represented by a :class:`tuple`. But if a tensor was treated as symbolic placeholder with tracing, it's shape could also be a :class:`~.Tensor`. See :class:`~.trace` for more details. The shape property is usually used to get the current shape of a tensor, but may also be used to reshape the tensor in-place by assigning a tuple of tensor dimensions to it. As with :func:`~.reshape`, one of the new shape dimensions can be -1, in which case its value is inferred from the size of the tensor and the remaining dimensions. """ shape = super().shape if shape == () or not use_symbolic_shape(): return shape return apply(GetVarShape(), self)[0] @property def _tuple_shape(self): return super().shape @property def device(self) -> CompNode: r"""Returns a string represents the device a :class:`~.Tensor` storaged on.""" return super().device @property def dtype(self) -> np.dtype: r"""Returns a :class:`numpy.dtype` object represents the data type of a :class:`~.Tensor`.""" return super().dtype @property def format(self) -> str: return super().format() @format.setter def format(self, format): super()._set_format(format) @property def qparams(self): r"""Returns a :class:`~.QParams` object containing quantization params of a :class:`~.Tensor`.""" from .quantization.utils import create_qparams # pylint: disable=all if self._qparams is None: self._qparams = create_qparams() return self._qparams
[文档] def numpy(self) -> np.ndarray: r"""Returns self :class:`~.Tensor` as a :class:`numpy.ndarray`.""" return super().numpy()
[文档] def detach(self): r"""Returns a new :class:`~.Tensor`, detached from the current graph.""" return super().detach()
def _reset(self, other): if not isinstance(other, _Tensor): other = Tensor(other, dtype=self.dtype, device=self.device) super()._reset(other) def __repr__(self): piece = "{}(".format(self.__class__.__name__) with np.printoptions(precision=4, suppress=True): piece += "{}".format(str(self.numpy())) if self.dtype != np.float32: piece += ", dtype={}".format(np.dtype(self.dtype).name) piece += ", device={}".format(self.device) + ")" return piece @property def name(self): return self._custom_name @name.setter def name(self, name): self._custom_name = name if name == None: name = "" self._name = self._prefix + "." + name if self._prefix else name self._set_name(self._name)
[文档] @deprecated( version="1.0", reason="please use ``tensor_name[...] = value``", ) def set_value(self, value): self._reset(value)
[文档] @deprecated(version="1.0", reason="use ``*= 0`` instead") def reset_zero(self): self *= 0
[文档] def to(self, device, *, _borrow=False): r"""Copy self :class:`~.Tensor` to specified device. See :func:`~.copy`""" if isinstance(device, str) and not _valid_device(device): raise ValueError( "invalid device name {}. For the correct format of the device name, please refer to the instruction of megengine.device.set_default_device()".format( device ) ) cn = as_device(device).to_c() op = Borrow(comp_node=cn) if _borrow else Copy(comp_node=cn) return apply(op, self)[0]
@property def requires_grad(self): raise AttributeError("requires_grad is reserved for future use") @requires_grad.setter def requires_grad(self, value): raise AttributeError("requires_grad is reserved for future use") @requires_grad.deleter def requires_grad(self): raise AttributeError("requires_grad is reserved for future use") def __hash__(self): return id(self) def __getnewargs__(self): r"""__getnewargs__ will be called for pickle serialization or deep copy""" return (self.numpy(), self.dtype, self.device.logical_name) def __getstate__(self): r"""__getstate__ will be called for pickle serialization or deep copy""" state = {} if self._qparams is not None: state["qparams"] = self._qparams return state def __setstate__(self, state): # for compatibility with old version not using fastcore if "data" in state: data = state.pop("data") device = state.pop("device") dtype = state.pop("dtype") self._reset(Tensor(data, dtype=dtype, device=device)) # quantize related state for deepcopy if "qdict" in state: qparams = state.pop("qdict") logger.warning( "Tensor's 'qdict' state is depreciated. Use 'qparams' instead" ) elif "qparams" in state: qparams = state.pop("qparams") else: qparams = None self._qparams = qparams
set_py_tensor_type(Tensor) tensor = Tensor
[文档]class Parameter(Tensor): r"""A kind of Tensor that is to be considered a module parameter. Note: Operations happened on Parameter usually return a Tensor instead of Parameter. For example, with a Parameter ``x``, ``x.reshape/to/sum/...`` will result into a Tensor. Any operations between Parameter and Tensor will have Tensor as outputs. """