提交 bfc4e7a9 编写于 作者: M Megvii Engine Team

docs(mge): fix amp docstring problems

GitOrigin-RevId: e5540eb9402c86d9d00f81de5b30ad94ce535eb6
上级 0b764cf2
...@@ -15,14 +15,14 @@ class autocast: ...@@ -15,14 +15,14 @@ class autocast:
A class to control autocast mode for amp as a context manager or a decorator. A class to control autocast mode for amp as a context manager or a decorator.
:param enabled: Whether autocast mode is enabled. :param enabled: Whether autocast mode is enabled.
:low_prec_dtype: Set amp autocast mode's lower precision dtype. It will change the :param low_prec_dtype: Set amp autocast mode's lower precision dtype. It will change
target dtype in tensor casting for better speed and memory. Default: float16. the target dtype in tensor casting for better speed and memory. Default: float16.
:high_prec_dtype: Set amp autocast mode's higher precision dtype. It will change the :param high_prec_dtype: Set amp autocast mode's higher precision dtype. It will
target dtype in tensor casting for better precision. Default: float32. change the target dtype in tensor casting for better precision. Default: float32.
Examples: Examples:
..code-block:: .. code-block::
# used as decorator # used as decorator
@autocast() @autocast()
......
...@@ -17,7 +17,7 @@ def enabled(mod): ...@@ -17,7 +17,7 @@ def enabled(mod):
Examples: Examples:
..code-block:: .. code-block::
import megengine as mge import megengine as mge
mge.amp.enabled = True mge.amp.enabled = True
...@@ -40,7 +40,7 @@ def high_prec_dtype(mod): ...@@ -40,7 +40,7 @@ def high_prec_dtype(mod):
Examples: Examples:
..code-block:: .. code-block::
import megengine as mge import megengine as mge
mge.amp.high_prec_dtype = "float32" mge.amp.high_prec_dtype = "float32"
...@@ -63,7 +63,7 @@ def low_prec_dtype(mod): ...@@ -63,7 +63,7 @@ def low_prec_dtype(mod):
Examples: Examples:
..code-block:: .. code-block::
import megengine as mge import megengine as mge
mge.amp.low_prec_dtype = "float16" mge.amp.low_prec_dtype = "float16"
......
...@@ -28,7 +28,7 @@ class QConfig( ...@@ -28,7 +28,7 @@ class QConfig(
) )
): ):
r""" r"""
A config class indicating how to do quantize toward :class:`~.QATModule`'s A config class indicating how to do quantize toward :class:`~.QATModule` 's
``activation`` and ``weight``. See :meth:`~.QATModule.set_qconfig` for detail usage. ``activation`` and ``weight``. See :meth:`~.QATModule.set_qconfig` for detail usage.
:param weight_observer: interface to instantiate an :class:`~.Observer` indicating :param weight_observer: interface to instantiate an :class:`~.Observer` indicating
......
...@@ -30,6 +30,7 @@ class Tensor(_Tensor, ArrayMethodMixin): ...@@ -30,6 +30,7 @@ class Tensor(_Tensor, ArrayMethodMixin):
A tensor object represents a multidimensional, homogeneous array of fixed-size items. A tensor object represents a multidimensional, homogeneous array of fixed-size items.
:param data: The value of returned Tensor. :param data: The value of returned Tensor.
:type data: Tensor, :class:`~.numpy.ndarray`, :class:`list` or python number.
:param dtype: The dtype of returned Tensor. Uses data's dtype if not specified. :param dtype: The dtype of returned Tensor. Uses data's dtype if not specified.
:param device: The desired device of returned Tensor. Uses :func:`get_default_device` if not specified. :param device: The desired device of returned Tensor. Uses :func:`get_default_device` if not specified.
:param is_const: Whether make it a ``ImutableTensor`` in tracing mode. :param is_const: Whether make it a ``ImutableTensor`` in tracing mode.
...@@ -43,7 +44,7 @@ class Tensor(_Tensor, ArrayMethodMixin): ...@@ -43,7 +44,7 @@ class Tensor(_Tensor, ArrayMethodMixin):
def __new__( def __new__(
cls, cls,
data: Union["Tensor", np.ndarray, list, "scalar"] = None, data: Union["Tensor", np.ndarray, list, int, float] = None,
dtype: np.dtype = None, dtype: np.dtype = None,
device: str = None, device: str = None,
is_const: bool = False, is_const: bool = False,
...@@ -76,7 +77,7 @@ class Tensor(_Tensor, ArrayMethodMixin): ...@@ -76,7 +77,7 @@ class Tensor(_Tensor, ArrayMethodMixin):
def __init__( def __init__(
self, self,
data: Union["Tensor", np.ndarray, list, "scalar"], data: Union["Tensor", np.ndarray, list, int, float],
dtype: np.dtype = None, dtype: np.dtype = None,
device: str = None, device: str = None,
is_const: bool = False, is_const: bool = False,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册