From 26a09fd582a7ecec8dc9100f7628505ee581f5ec Mon Sep 17 00:00:00 2001 From: bjutsecurity22 Date: Fri, 1 Sep 2023 16:18:40 +0800 Subject: [PATCH] Update cell.py --- mindspore/nn/cell.py | 2033 ++++++++++++++++++++++++++++++++++++------ 1 file changed, 1738 insertions(+), 295 deletions(-) diff --git a/mindspore/nn/cell.py b/mindspore/nn/cell.py index 50430c8cc4..8b5d5bf832 100755 --- a/mindspore/nn/cell.py +++ b/mindspore/nn/cell.py @@ -1,4 +1,4 @@ -# Copyright 2020 Huawei Technologies Co., Ltd +# Copyright 2020-2022 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,43 +18,43 @@ import inspect import os import time from collections import OrderedDict - +from types import FunctionType, MethodType import numpy +from mindspore._checkparam import args_type_check from mindspore import log as logger from mindspore.common.parameter import PARAMETER_NAME_DEFAULT -from mindspore.common._decorator import deprecated +from mindspore.common.hook_handle import HookHandle from mindspore.context import ParallelMode +from mindspore.ops.composite import Shard from .. import context -from .._c_expression import init_pipeline, Cell_, FuncGraph +from .._c_expression import init_pipeline, update_func_graph_hyper_params, Cell_, FuncGraph, MixedPrecisionType from .._checkparam import Validator from ..common import dtype as mstype -from ..common.api import _executor, _pynative_exec +from ..common.api import _cell_graph_executor, _pynative_executor, _check_all_tensor, cells_compile_cache from ..common.parameter import Parameter, ParameterTuple -from ..common.tensor import Tensor -from ..ops.functional import cast -from ..ops.operations import HookBackward +from ..common.variable import Variable +from ..common.tensor import Tensor, CSRTensor, COOTensor +from ..ops.operations import Cast from ..ops.primitive import Primitive +from ..ops.operations import _inner_ops as inner from ..parallel._tensor import _load_tensor_by_layout class Cell(Cell_): """ - Base class for all neural networks. - - A 'Cell' could be a single neural network cell, such as conv2d, relu, batch_norm, etc. or a composition of - cells to constructing a network. + The basic building block of neural networks in MindSpore. The model or neural network layer should inherit this + base class. - Note: - In general, the autograd algorithm will automatically generate the implementation of the gradient function, - but if back-propagation(bprop) method is implemented, the gradient function will be replaced by the bprop. - The bprop implementation will receive a Tensor `dout` containing the gradient of the loss w.r.t. - the output, and a Tensor `out` containing the forward result. The bprop needs to compute the - gradient of the loss w.r.t. the inputs, gradient of the loss w.r.t. Parameter variables are not supported - currently. The bprop method must contain the self parameter. + Layers in `mindspore.nn` are also the subclass of Cell, such as :class:`mindspore.nn.Conv2d`, + :class:`mindspore.nn.ReLU`, :class:`mindspore.nn.BatchNorm`, etc. Cell will be compiled into a calculation + graph in GRAPH_MODE (static graph mode) and used as the basic module of neural networks in + PYNATIVE_MODE (dynamic graph mode). Args: - auto_prefix (bool): Recursively generate namespaces. Default: True. + auto_prefix (bool): Whether to automatically generate NameSpace for Cell and its subcells. It will affect the + name of the parameter in the network. If set to True, the network parameter + name will be prefixed, otherwise it will not. Default: True. flags (dict): Network configuration information, currently it is used for the binding of network and dataset. Users can also customize network attributes by this parameter. Default: None. @@ -62,41 +62,94 @@ class Cell(Cell_): ``Ascend`` ``GPU`` ``CPU`` Examples: + >>> import mindspore.nn as nn + >>> import mindspore.ops as ops >>> class MyCell(nn.Cell): - ... def __init__(self): - ... super(MyCell, self).__init__() - ... self.relu = P.ReLU() + ... def __init__(self, forward_net): + ... super(MyCell, self).__init__(auto_prefix=False) + ... self.net = forward_net + ... self.relu = ops.ReLU() ... - ... def construct(self, x): - ... return self.relu(x) + ... def construct(self, x): + ... y = self.net(x) + ... return self.relu(y) + >>> + >>> inner_net = nn.Conv2d(120, 240, 4, has_bias=False, weight_init='normal') + >>> my_net = MyCell(inner_net) + >>> print(my_net.trainable_params()) + ... # If the 'auto_prefix' set to True or not set when call the '__init__' method of the parent class, + ... # the parameter's name will be 'net.weight'. + [Parameter (name=weight, shape=(240, 120, 4, 4), dtype=Float32, requires_grad=True)] """ + + class _CellGuard: + """Detecting whether the cell is a top-level cell with the 'with statement'.""" + def __enter__(self):#输入cell并增加递归深度计数 + """Enter cell and increase recursion depth count.""" + _pynative_executor.set_lazy_build(True) + _pynative_executor.enter_cell() + + def __exit__(self, exc_type, exc_val, exc_tb):#退出cell并减少递归深度计数 + """Exit cell and decrease recursion depth count.""" + _pynative_executor.exit_cell() + if _pynative_executor.is_top_cell(): + _pynative_executor.set_lazy_build(False) + IGNORE_LIST = ['_scope', '_cell_init_args', '_auto_prefix', '_cells', '_params', '_construct_inputs_names', '_construct_inputs_num', '_create_time', '_mindspore_flags', '_parallel_inputs_run', - '_parameter_layout_dict', '_params_list', '_tensor_list', '_phase', - '_auto_parallel_mode', '_backward_hook', '_bprop_debug', '_is_run', '_param_prefix', - '_attr_synced', 'enable_hook', 'pynative', 'requires_grad', - '_auto_parallel_compile_and_run', 'cell_type'] + '_parameter_layout_dict', '_params_list', '_tensor_list', '_phase', '_auto_parallel_mode', + '_forward_pre_hook', '_forward_hook', '_enable_forward_pre_hook', '_enable_forward_hook', + '_bprop_debug', '_enable_backward_hook', '_cell_backward_hook', '_is_run', '_param_prefix', + '_attr_synced', 'pynative', 'requires_grad', '_auto_parallel_compile_and_run', 'cell_type'] def __init__(self, auto_prefix=True, flags=None): + # 初始化Cell类 Cell_.__init__(self, self._cell_tag) + # 初始化参数字典 self._params = OrderedDict() + # 初始化单元字典 self._cells = OrderedDict() + # 初始化参数列表 self._params_list = OrderedDict() + # 初始化矢量列表 self._tensor_list = OrderedDict() + # 初始化操作字典 + self._primitives = OrderedDict() + # 设置是否自动添加前缀 self.training = False + # 设置是否需要梯度 self.requires_grad = False + # 设置是否使用pynative self.pynative = False + # 设置参数前缀 self._attr_synced = False + # 设置参数前缀 self._param_prefix = '' + # 设置是否自动添加前缀 self._auto_prefix = auto_prefix + # 设置作用域 self._scope = None + # 设置训练阶段 self._phase = 'train' + # 设置参数排序字典 self._parameter_layout_dict = {} + # 设置并行参数名称列表 self._parallel_parameter_name_list = () + # 设置并行参数合并网络字典 self._parallel_parameter_merge_net_dict = {} + # 创建时间 self._create_time = int(time.time() * 1e9) - self.phase_prefix = "" + # 初始化参数 + self.arguments_key = "" + # 初始化编译缓存 + self.compile_cache = set() + # 初始化单元编译缓存 + cells_compile_cache[id(self)] = self.compile_cache + # 设置参数广播是否完成 self.parameter_broadcast_done = False + self._id = 1 + self.exist_names = set("") + self.exist_objs = set() init_pipeline() # call gc to release GE session resources used by non-used cell objects @@ -105,24 +158,54 @@ class Cell(Cell_): self._construct_inputs_num = 0 self._construct_inputs_names = [] + # 设置是否使用自动平行模式 self._auto_parallel_mode = False + # 设置并行参数运行 self._parallel_inputs_run = None + # 初始化前向调用钩子 if flags: self.add_flags(**flags) - self._backward_hook = None - self.enable_hook = False + # 设置是否启用前向调用钩子 self._bprop_debug = False + # 初始化前向前置钩子 + self._forward_pre_hook = OrderedDict() + # 初始化前向钩子 + self._forward_hook = OrderedDict() + # 初始化是否启用前向前置钩子 + self._enable_forward_pre_hook = False + # 初始化是否启用前向钩子 + self._enable_forward_hook = False + # 初始化是否启用后向钩子 + self._enable_backward_hook = False + # 初始化单元后向钩子 + self._cell_backward_hook = None + # 设置单元类型 self.cell_type = None + # 设置自动平行编译和运行 self._auto_parallel_compile_and_run = False + # 初始化参数 + self.cast = Cast() + # 初始化是否有配置重新计算 + self._has_config_recompute = False + # 初始化用户参数 + self._user_parameters = [] + self._dynamic_shape_inputs = None + self.saved_dynamic_shape = None def __getstate__(self): + # 获取Cell类的状态 base = Cell_.__getstate__(self) + # 返回base和self.__dict__ return base, self.__dict__ def __setstate__(self, state): + # 获取state base, dict_ = state + # 设置Cell类的状态 Cell_.__setstate__(self, base) + # 将self.__dict__设置为dict_ self.__dict__ = dict_ + # 将_attr_synced设置为False self._attr_synced = False @property @@ -146,14 +229,14 @@ class Cell(Cell_): return self._param_prefix @property - def bprop_debug(self): + def bprop_debug(self):#获取是否启用了cell自定义bprop调试 """ Get whether cell custom bprop debug is enabled. """ return self._bprop_debug @bprop_debug.setter - def bprop_debug(self, value): + def bprop_debug(self, value):#设置是否启用cell自定义bprop调试 """ Set whether to enable cell custom bprop debug. @@ -166,10 +249,10 @@ class Cell(Cell_): value (bool): Specifies whether to enable bprop debug. Default: False. """ if not isinstance(value, bool): - raise TypeError("'bprop debug' value must be bool type.") + raise TypeError(f"For 'Cell', the property 'bprop_debug' must be bool type, but got type {type(value)}.") self._bprop_debug = value - def update_cell_prefix(self): + def update_cell_prefix(self):#更新所有child cells' self.param_prefix. """ Update the all child cells' self.param_prefix. @@ -180,73 +263,135 @@ class Cell(Cell_): for cell_name, cell in cells_name: cell._param_prefix = cell_name - def update_cell_type(self, cell_type): + def update_cell_type(self, cell_type):#更新 cell 格式为 'cell_type' """ The current cell type is updated when a quantization aware training network is encountered. After being invoked, it can set the cell type to 'cell_type'. + + Args: + cell_type(str): The type of cell to be updated, cell_type can be "quant" or "second-order". """ self.cell_type = cell_type @cell_init_args.setter def cell_init_args(self, value): + ''' + 设置cell的初始化参数 + ''' if not isinstance(value, str): - raise TypeError("'cell_init_args' must be string type.") + raise TypeError(f"For 'Cell', the property 'cell_init_args' must be string type, " + f"but got type {type(value)}.") self._cell_init_args = value @property def phase(self): + # 获取Cell的阶段 return self._phase @phase.setter def phase(self, value): + # 设置Cell的阶段 if not isinstance(value, str): - raise TypeError("'phase' must be string type.") + raise TypeError(f"For 'Cell', the property 'phase' must be string type, but got type {type(value)}.") self._phase = value @property def parameter_layout_dict(self): + """ + `parameter_layout_dict` represents the tensor layout of a parameter, which is inferred by shard strategy and + distributed operator information. + """ + # 获取Cell的参数布局字典 return self._parameter_layout_dict @property def cls_name(self): + # 获取Cell的类名 return self.__class__.__name__ @parameter_layout_dict.setter def parameter_layout_dict(self, value): + # 设置Cell的参数布局字典 if not isinstance(value, dict): - raise TypeError("'parameter_layout_dict' must be dict type.") + raise TypeError(f"For 'Cell', the property 'parameter_layout_dict' must be dict type, " + f"but got type {type(value)}.") self._parameter_layout_dict = value @property def parallel_parameter_name_list(self): + # 获取Cell的并行参数名称列表 return self._parallel_parameter_name_list @parallel_parameter_name_list.setter def parallel_parameter_name_list(self, value): + ''' + 设置参数名称列表 + ''' if not isinstance(value, list): - raise TypeError("'parallel_parameter_name_list' must be list type.") + raise TypeError(f"For 'Cell', the property 'parallel_parameter_name_list' must be list type, " + f"but got type {type(value)}.") self._parallel_parameter_name_list = value + @property + def pipeline_stage(self): + return self._pipeline_stage + + @pipeline_stage.setter + def pipeline_stage(self, value): + ''' + 设置pipeline_stage属性 + :param value: int类型，表示pipeline_stage的值 + :return: None + ''' + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError("For 'Cell', the property 'pipeline_stage' " + "must be int type, but got type : {}".format(type(value))) + + if value < 0: + raise ValueError("For 'Cell', the property 'pipeline_stage' " + "can not be less than 0, but got {}".format(value)) + self._pipeline_stage = value + for item in self.trainable_params(): + item.add_pipeline_stage(value) + @property def parallel_parameter_merge_net_dict(self): + ''' + 获取parallel_parameter_merge_net_dict属性 + :return: dict类型 + ''' return self._parallel_parameter_merge_net_dict @parallel_parameter_merge_net_dict.setter def parallel_parameter_merge_net_dict(self, value): + ''' + 设置parallel_parameter_merge_net_dict属性 + :param value: dict类型 + :return: None + ''' if not isinstance(value, dict): - raise TypeError("'parallel_parameter_merge_net_dict' must be dict type.") + raise TypeError(f"For 'Cell', the property 'parallel_parameter_merge_net_dict' must be dict type, " + f"but got type {type(value)}.") self._parallel_parameter_merge_net_dict = value def get_func_graph_proto(self): """Return graph binary proto.""" - return _executor._get_func_graph_proto(self, self.phase + "." + str(self.create_time), "anf_ir", True) + ''' + 返回graph binary proto + :return: graph binary proto + ''' + exec_id = ".".join([self.phase, str(self.create_time), str(id(self))]) + return _cell_graph_executor._get_func_graph_proto(self, exec_id, "anf_ir", True) def __getattr__(self, name): + ''' + 获取属性 + ''' if '_params' in self.__dict__: params = self.__dict__['_params'] if name in params: - if context.get_context("mode") == context.PYNATIVE_MODE: + if context._get_mode() == context.PYNATIVE_MODE: return self.cast_param(params[name]) return params[name] if '_cells' in self.__dict__: @@ -266,15 +411,25 @@ class Cell(Cell_): cast_list.append(self.cast_param(para)) para_list = ParameterTuple(cast_list) return para_list - raise AttributeError("'{}' object has no attribute '{}'.".format(type(self).__name__, name)) + raise AttributeError("The '{}' object has no attribute '{}'.".format(type(self).__name__, name)) def __del__(self): - if context.get_context is not None and context.get_context("mode") == context.PYNATIVE_MODE: - _pynative_exec.del_cell(str(id(self))) - if hasattr(self, "_create_time"): - _executor.del_net_res(str(self._create_time)) + ''' + 析构函数 + ''' + if context.get_context is not None and context._get_mode() == context.PYNATIVE_MODE: + _pynative_executor.del_cell(str(id(self))) + + # while deepcopy a cell instance, the copied cell instance can't be added to cells_compile_cache + # here using pop(id(self), None) to avoid KeyError exception + cells_compile_cache.pop(id(self), None) + if self.compile_cache: + _cell_graph_executor.del_net_res(self.compile_cache) def __delattr__(self, name): + ''' + 删除属性 + ''' if name in self._params: del self._params[name] elif name in self._cells: @@ -291,115 +446,365 @@ class Cell(Cell_): """Cast input for mixed precision""" res = list() for item in inputs: + # 如果item是元组类型 if isinstance(item, tuple): + # 将元组中的元素转换为指定类型 res.append(self._cast_mixed_precision_inputs(item, dst_type)) + # 如果item是浮点数 elif isinstance(item, float): - res.append(cast(item, dst_type)) + # 将元组中的元素转换为指定类型 + res.append(self.cast(item, dst_type)) + # 如果item有dtype属性，且dtype属性的值在{mstype.float16, mstype.float32, mstype.float64}中 elif hasattr(item, "dtype") and item.dtype in {mstype.float16, mstype.float32, mstype.float64}: - res.append(cast(item, dst_type)) + # 将元组中的元素转换为指定类型 + res.append(self.cast(item, dst_type)) + # 否则 else: + # 将元组中的元素转换为指定类型 res.append(item) + # 返回转换后的元组 return tuple(res) def cast_inputs(self, inputs, dst_type): + """ + Cast inputs to specified type. + + Args: + inputs (tuple[Tensor]): The cell inputs. + dst_type (mindspore.dtype): The specified data type. + + returns: + tuple[Tensor], the result with destination data type. + """ res = list() for item in inputs: + # 如果item是元组类型 if isinstance(item, tuple): + # 调用cast_inputs函数，将元组转换为dst_type类型 res.append(self.cast_inputs(item, dst_type)) + # 否则 else: - res.append(cast(item, dst_type)) + # 调用cast函数，将item转换为dst_type类型 + res.append(self.cast(item, dst_type)) + # 返回转换后的结果 return tuple(res) - def do_parameter_broadcast(self): + def _do_parameter_broadcast(self): if context.get_auto_parallel_context("parallel_mode") == ParallelMode.DATA_PARALLEL: if not self.parameter_broadcast_done: - _pynative_exec.parameter_broadcast(self, self.phase, self._auto_parallel_mode) + _pynative_executor.parameter_broadcast(self, self.phase, self._auto_parallel_mode) self.parameter_broadcast_done = True def run_construct(self, cast_inputs, kwargs): - if self.enable_hook: - _pynative_exec.enter_construct(self) - output = self._hook_construct(*cast_inputs, **kwargs) - _pynative_exec.leave_construct(self) + """ + Run the construct function. + + Note: + This function will be removed in a future version. It is not recommended to call this function. + + Args: + cast_inputs (tuple): The input objects of Cell. + kwargs (dict): Provide keyword arguments. + + Returns: + output, the output object of Cell. + """ + logger.warning(f"The 'run_construct' function of '{self.cls_name}' will be removed in a future version. " + f"Calling this function is not recommended.") + output = self._run_construct(cast_inputs, kwargs) + return output + + def _run_construct(self, cast_inputs, kwargs): + """Run the construct function""" + if self._enable_forward_pre_hook: + # 调用_run_forward_pre_hook函数 + cast_inputs = self._run_forward_pre_hook(cast_inputs) + # 如果_enable_backward_hook为True，则调用_backward_hook_construct函数 + if self._enable_backward_hook: + output = self._backward_hook_construct(*cast_inputs) + # 如果_enable_backward_hook为False，且_shard_fn存在，则调用_shard_fn函数 + elif hasattr(self, "_shard_fn"): + output = self._shard_fn(*cast_inputs, **kwargs) + # 否则，调用self.construct函数 else: - _pynative_exec.enter_construct(self) output = self.construct(*cast_inputs, **kwargs) - _pynative_exec.leave_construct(self) + # 如果_enable_forward_hook为True，则调用_run_forward_hook函数 + if self._enable_forward_hook: + output = self._run_forward_hook(cast_inputs, output) + # 返回输出 return output - def __call__(self, *inputs, **kwargs): + def _check_construct_args(self, *inputs, **kwargs): + """Check the args needed by the function construct""" + if kwargs: + raise ValueError(f"For 'Cell', expect no kwargs here, " + "maybe you pass wrong arguments, args: {inputs}, kwargs: {kwargs}") + # 初始化positional_args和default_args + positional_args = 0 + default_args = 0 + # 遍历构造函数的参数 + for value in inspect.signature(self.construct).parameters.values(): + # 如果参数是可变参数或者可变关键字参数 + if value.kind is inspect.Parameter.VAR_POSITIONAL or value.kind is inspect.Parameter.VAR_KEYWORD: + return + # 跳过 + continue + # 如果参数是正常参数 + if value.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD: + # 如果参数默认值为空 + if value.default is inspect.Parameter.empty: + # 将positional_args加1 + positional_args += 1 + else: + # 将default_args加1 + default_args += 1 + + # 如果输入的参数小于positional_args + if len(inputs) < positional_args: + # 抛出TypeError + raise TypeError(f"For 'Cell', the function construct need {positional_args} positional argument, " + f"but got {len(inputs)}.") + + # 如果输入的参数大于positional_args + default_args + if len(inputs) > positional_args + default_args: + # 抛出TypeError + raise TypeError(f"For 'Cell', the function construct need {positional_args} positional argument and " + f"{default_args} default argument, total {positional_args + default_args}, " + f"but got {len(inputs)}.") + + def _hook_fn_registered(self): + if self._enable_forward_pre_hook or self._enable_forward_hook or self._enable_backward_hook: + return True + # 遍历cells + for cell in self.cells(): + # 如果cell没有注册hook_fn + if cell._hook_fn_registered(): + return True + return False + + def _get_prims_recursively(self): + all_prims = list() + # 遍历_primitives字典，将每一个值转换为元组，并将元组添加到all_prims列表中 + for _, value in self._primitives.items(): + if value: + all_prims.append(value) + + # 遍历cells，将每一个cell的_get_prims_recursively()方法的返回值添加到all_prims列表中 + for cell in self.cells(): + all_prims.extend(cell._get_prims_recursively()) + + return all_prims + + def set_data_parallel(self):#并行化所有ops算子 + """ + For all primitive ops in this cell(including ops of cells that wrapped by this cell), + if parallel strategy is not specified, then instead of auto-searching, data parallel + strategy will be generated for those primitive ops. + + Note: + Only effective while using auto_parallel_context = ParallelMode.AUTO_PARALLEL under graph mode. + + Examples: + >>> import mindspore.nn as nn + >>> net = nn.Dense(3, 4) + >>> net.set_data_parallel() + """ + if context._get_mode() == context.PYNATIVE_MODE: + raise ValueError("set_data_parallel: does not support PyNative mode.") + + all_prims = self._get_prims_recursively() + for prim in all_prims: + prim.add_prim_attr("strategy_gen_mode", "data_parallel") + + def shard(self, in_strategy, out_strategy, device="Ascend", level=0): + """ + Defining the input and output layouts of this cell and the parallel strategies of remaining ops will be + generated by sharding propagation. in_strategy and out_strategy define the input and output layout respectively. + in_strategy/out_strategy should be a tuple each element of which corresponds to the desired layout of + this input/output and None represents data_parallel. + + Note: + Only effective in PYNATIVE_MODE and in either ParallelMode.AUTO_PARALLEL with + search_mode in auto_parallel_context set as sharding_propagation. + + Args: + in_strategy (tuple): Define the layout of inputs, each element of the tuple should be a tuple or None. Tuple + defines the layout of the corresponding input and None represents a data parallel strategy. + out_strategy (tuple): Define the layout of outputs similar with in_strategy. + device (string): Select a certain device target. It is not in use right now. + Support ["CPU", "GPU", "Ascend"]. Default: "Ascend". + level (int): Option for parallel strategy infer algorithm, namely the object function, maximize computation + over communication ratio, maximize speed performance, minimize memory usage etc. It is not in + use right now. Support ["0", "1", "2"]. Default: "0". + + Returns: + Cell, the cell itself. + + Examples: + >>> import mindspore.nn as nn + >>> + >>> class Block(nn.Cell): + ... def __init__(self): + ... self.dense1 = nn.Dense(10, 10) + ... self.relu = nn.ReLU() + ... self.dense2 = nn.Dense2(10, 10) + ... def construct(self, x): + ... x = self.relu(self.dense2(self.relu(self.dense1(x)))) + ... return x + >>> + >>> class example(nn.Cell): + ... def __init__(self): + ... self.block1 = Block() + ... self.block2 = Block() + ... self.block2.shard(in_strategy=((2, 1),), out_strategy=(None,)) + ... def construct(self, x): + ... x = self.block1(x) + ... x = self.block2(x) + ... return x + """ + shard_fn = Shard() + # 创建一个Shard函数，用于分割输入和输出 + fn = shard_fn(self, in_strategy, out_strategy, device, level) + # 将分割函数赋值给object.__setattr__(self, "_shard_fn", fn) + object.__setattr__(self, "_shard_fn", fn) + # 返回self + return self + + def auto_cast_inputs(self, inputs): + """ + Auto cast inputs in mixed precision scenarios. + + Args: + inputs (tuple): the inputs of construct. + + Returns: + Tuple, the inputs after data type cast. + """ + cast_inputs = inputs + # 获取当前类型的调整精度类型 + mixed_type = self.get_mixed_precision_type() + # 如果当前类型的调整精度类型为FP16，则将输入转换为FP16类型 + if mixed_type == MixedPrecisionType.FP16: + cast_inputs = self._cast_mixed_precision_inputs(inputs, mstype.float16) + # 如果当前类型的调整精度类型为FP32，则将输入转换为FP32类型 + if mixed_type == MixedPrecisionType.FP32: + cast_inputs = self._cast_mixed_precision_inputs(inputs, mstype.float32) + + return cast_inputs + + def __call__(self, *args, **kwargs): if self.__class__.construct is Cell.construct: + # 如果类的构造函数没有被重写，则警告日志 logger.warning(f"The '{self.__class__}' does not override the method 'construct', " - f"will call the super class(Cell) 'construct'.") + f"it will call the super class(Cell) 'construct'.") + # 如果有参数，则使用inspect模块的signature函数绑定参数 if kwargs: - bound_args = inspect.signature(self.construct).bind(*inputs, **kwargs) - inputs = bound_args.args - kwargs = bound_args.kwargs - if context.get_context("mode") == context.GRAPH_MODE: - if kwargs: - raise ValueError("For 'graph' mode, the outermost network does not support passing " - "variable key-value pair parameters.") - if self.enable_hook: - raise ValueError("The graph mode does not support hook function.") - out = self.compile_and_run(*inputs) + bound_arguments = inspect.signature(self.construct).bind(*args, **kwargs) + bound_arguments.apply_defaults() + args = bound_arguments.args + kwargs = bound_arguments.kwargs + + # Run in Graph mode. + if context._get_mode() == context.GRAPH_MODE: + # 检查构造参数 + self._check_construct_args(*args, **kwargs) + # 检查hook函数是否已经注册 + if self._hook_fn_registered(): + logger.warning(f"For 'Cell', it's not support hook function in graph mode. If you want to use hook " + f"function, please use context.set_context to set pynative mode.") + # 编译和运行 + out = self.compile_and_run(*args) return out - self.do_parameter_broadcast() - for item in inputs: - if isinstance(item, numpy.ndarray): - raise TypeError("cell inputs should not be numpy array.") - origin_grad = [] - if self.requires_grad is True: - _pynative_exec.set_grad_flag(True) - _pynative_exec.new_graph(self, *inputs, **kwargs) - for cell in self.cells(): - origin_grad.append(cell.requires_grad) - cell.set_grad(True) - else: - _pynative_exec.set_grad_flag(False) - cast_inputs = list() - if hasattr(self, "_mindspore_flags"): - if self._mindspore_flags.get('fp16'): - cast_inputs = self._cast_mixed_precision_inputs(inputs, mstype.float16) - if self._mindspore_flags.get('fp32'): - cast_inputs = self._cast_mixed_precision_inputs(inputs, mstype.float32) - if not cast_inputs: - cast_inputs = inputs - output = self.run_construct(cast_inputs, kwargs) + # Run in PyNative mode. + if _pynative_executor.is_top_cell(): + # 如果是最顶层的cell，则设置lazy_build为True，并设置optimizer为None + _pynative_executor.set_lazy_build(True) + _pynative_executor._optimizer = getattr(self, "optimizer", None) + _pynative_executor._top_cell = self + # There many Casts in parameter_broadcast. Enable lazy_build and build faster. + # 设置参数broadcast的状态 + self._do_parameter_broadcast() + + # 将参数转换为numpy数组 + for item in args: + if isinstance(item, Tensor) and item.has_init: + item.init_data() + elif isinstance(item, numpy.ndarray): + raise TypeError("For 'Cell', inputs should not be numpy array.") + # 如果需要梯度，则设置梯度标志为True + if self.requires_grad: + _pynative_executor.set_grad_flag(True) + # 创建新图 + _pynative_executor.new_graph(self, *args, **kwargs) + # 获取自动转换的输入 + cast_inputs = self.auto_cast_inputs(args) + + with self._CellGuard(): + try: + # 运行构造函数 + output = self._run_construct(cast_inputs, kwargs) + except Exception as err: + # 清除资源 + _pynative_executor.clear_res() + raise err + + if _pynative_executor.is_top_cell(): + # 如果是最顶层的单元，则执行lazy_task + _pynative_executor.execute_lazy_task() + + # 如果output是Parameter类型，则将output转换为data if isinstance(output, Parameter): output = output.data - if self.requires_grad is True: - _pynative_exec.end_graph(self, output, *inputs, **kwargs) - for i, cell in enumerate(self.cells()): - cell.set_grad(origin_grad[i]) + # 结束图，并将output传入_pynative_executor.end_graph + _pynative_executor.end_graph(self, output, *args, **kwargs) + # 返回output return output - def _add_attr(self, name, value): - if name and name[:2] != '__' and name not in Cell.IGNORE_LIST: + ''' + 添加属性 + :param name: 属性名 + :param value: 属性值 + :return: None + ''' + if name and name[:2]!= '__' and name not in Cell.IGNORE_LIST: super(Cell, self)._add_attr(name, value) def _sync_attr_for_compile(self): """Sync the attr to c++ object.""" if self._attr_synced: return + # 获取cells cells = self.__dict__.get('_cells') + # 遍历cells for key in cells: cell = cells[key] + # 调用cell的_sync_attr_for_compile()方法 cell._sync_attr_for_compile() + # 添加attr self._add_attr(key, cell) + # 获取params params = self.__dict__.get('_params') + # 遍历params for key in params: if '.' in key: continue param = params[key] + # 添加attr self._add_attr(key, param) + # 获取params_list params_list = self.__dict__.get('_params_list') + # 遍历params_list for key in params_list: params_list_item = params_list[key] + # 添加attr self._add_attr(key, params_list_item) + # 遍历self.__dict__ for key in self.__dict__: value = self.__dict__[key] + # 添加attr self._add_attr(key, value) + # 将attr同步状态设置为True self._attr_synced = True def _set_attr_for_parameter(self, name, value): @@ -407,87 +812,190 @@ class Cell(Cell_): cells = self.__dict__.get('_cells') params = self.__dict__.get('_params') if params is None: - raise AttributeError("Can not assign params before Cell.__init__() call.") + raise AttributeError("For 'Cell', can not assign params before Cell.__init__() is called.") if name in self.__dict__: if self.__dict__[name] is not None: - raise TypeError("The type of value should not be Parameter or Cell, but got Parameter.") + raise TypeError(f"For 'Cell', the {name} should not be Parameter.") del self.__dict__[name] if cells and name in cells: - raise TypeError("The type of value should be Cell, but got Parameter.") + raise TypeError(f"For 'Cell', the {name} should be Cell, but got Parameter.") self.insert_param_to_cell(name, value) def _set_attr_for_parameter_tuple(self, name, value): - """Set attr for parameter tuple.""" + """Set attr for parameter in ParameterTuple.""" + # 获取Cell的参数 params = self.__dict__.get('_params') + # 获取Cell的参数列表 params_list = self.__dict__.get('_params_list') + # 如果没有参数，则抛出异常 if params is None: - raise AttributeError("Can not assign params before Cell.__init__() call.") + raise AttributeError("For 'Cell', can not assign params before Cell.__init__() is called.") + # 初始化已存在的名称 + exist_names = set("") + exist_objs = set() + # 遍历value for item in value: - self.insert_param_to_cell(item.name, item, check_name=False) - if context.get_context("mode") == context.PYNATIVE_MODE: + # 如果有多个相同的对象，它们的名称只检查一次 + if item in exist_objs: + # If there are multiple identical objects, their names only check once. + continue + # 将item添加到已存在的名称中 + exist_objs.add(item) + # 如果item的名称为默认名称，则警告 + if item.name == PARAMETER_NAME_DEFAULT: + logger.warning("For 'Cell', the parameter definition is deprecated.\n" + "Please set a unique name for the parameter in ParameterTuple '{}'.".format(value)) + # 将item的名称添加到Cell的参数列表中 + item.name = item.name + "$" + str(self._id) + self._id += 1 + # 添加参数到Cell的参数列表 + self.insert_param_to_cell(item.name, item, check_name_contain_dot=False) + # 如果已存在的名称中包含exist_names，则抛出异常 + if item.name in exist_names: + raise ValueError("The value {}, its name '{}' already exists. " + "Please set a unique name for the parameter.".format(value, item.name)) + # 将已存在的名称添加到已存在的名称中 + exist_names.add(item.name) + + if context._get_mode() == context.PYNATIVE_MODE: + # 如果当前模式为PYNATIVE_MODE if name in self.__dict__: + # 如果name在self.__dict__中 del self.__dict__[name] + # 删除self.__dict__中name if name in params: + # 如果name在params中 del params[name] + # 删除params中name params_list[name] = value + # 将value添加到params_list中 else: + # 如果当前模式为NATIVE_MODE object.__setattr__(self, name, value) + def _set_attr_for_parameter_in_list_or_tuple(self, name, value): + """Set attr for parameter in list or tuple.""" + for item in value: + # 如果value中存在item，则跳过本次循环 + if item in self.exist_objs: + # If there are multiple identical objects, their names only check once. + # 如果value中存在重复的对象，则仅检查一次 + continue + self.exist_objs.add(item) + # 将item添加到exist_objs中 + if item.name == PARAMETER_NAME_DEFAULT: + # 如果item的name为默认值，则将item的name添加到item的name中 + item.name = item.name + "$" + str(self._id) + self._id += 1 + # 如果item的name已经存在于exist_names中，则抛出异常 + if item.name in self.exist_names: + raise ValueError("The value {}, its name '{}' already exists. " + "Please set a unique name for the parameter.".format(value, item.name)) + # 如果item的name已经存在于exist_names中，则抛出异常 + self.exist_names.add(item.name) + # 将item的name添加到exist_names中 + object.__setattr__(self, name, value) + def _set_attr_for_cell(self, name, value): """Set attr for cell.""" cells = self.__dict__.get('_cells') params = self.__dict__.get('_params') if cells is None: - raise AttributeError("Can not assign cells before Cell.__init__() call.") + raise AttributeError("For 'Cell', can not assign cells before Cell.__init__() is called.") if name in self.__dict__: del self.__dict__[name] if params and name in params: - raise TypeError("The type of value should be Parameter, but got Cell.") + raise TypeError(f"For 'Cell', the {name} should be Parameter, but got Cell.") if self._auto_prefix: + # 更新参数名称 value.update_parameters_name(name + '.') + # 将参数添加到cells字典中 cells[name] = value if hasattr(self, '_cell_init_args'): + # 如果存在_cell_init_args属性 self.cell_init_args += str({name: value}) + def _set_attr_for_params(self, name, value): + if isinstance(value, Tensor) and self._params[name] is not None: + # 如果value是Tensor类型，且name对应的参数不为空 + self._params[name].set_data(value) + elif value is not None: + # 如果value不为空 + raise TypeError(f"For 'Cell', the type of {name} should be Parameter or ParameterTuple, " + f"but got {type(value).__name__}.") + else: + # 如果value为空 + self.insert_param_to_cell(name, None) + + def _set_attr_for_tensor(self, name, value): + if context._get_mode() == context.PYNATIVE_MODE: + # 如果当前模式为PYNATIVE_MODE，则将tensor_list设置为self.__dict__.get('_tensor_list') + tensor_list = self.__dict__.get('_tensor_list') + # 获取self.__dict__中的name + if name in self.__dict__: + # 如果name存在，则删除self.__dict__中的name + del self.__dict__[name] + # 将value添加到tensor_list中 + tensor_list[name] = value + else: + # 如果当前模式为NATIVE_MODE，则将self.__dict__[name]设置为value + object.__setattr__(self, name, value) + def __setattr__(self, name, value): + # 获取cells属性 cells = self.__dict__.get('_cells') + # 获取params属性 params = self.__dict__.get('_params') - tensor_list = self.__dict__.get('_tensor_list') + # 如果value是Parameter类型 if isinstance(value, Parameter): + # 调用_set_attr_for_parameter方法 self._set_attr_for_parameter(name, value) + # 如果value是ParameterTuple类型 elif isinstance(value, ParameterTuple): + # 调用_set_attr_for_parameter_tuple方法 self._set_attr_for_parameter_tuple(name, value) + # 如果value是list或tuple类型，且value为可检查的参数列表或元组 + elif isinstance(value, (list, tuple)) and value and _check_param_list_tuple(value): + # 调用_set_attr_for_parameter_in_list_or_tuple方法 + self._set_attr_for_parameter_in_list_or_tuple(name, value) + # 如果value是Cell类型 elif isinstance(value, Cell): + # 调用_set_attr_for_cell方法 self._set_attr_for_cell(name, value) + # 如果params属性为真 elif params and name in params: - if isinstance(value, Tensor) and self._params[name] is not None: - self._params[name].set_data(value) - elif value is not None: - raise TypeError(f"The type of value should be Parameter or ParameterTuple, " - f"but got {type(value).__name__}.") - else: - self.insert_param_to_cell(name, None) + # 调用_set_attr_for_params方法 + self._set_attr_for_params(name, value) + # 如果cells属性为真 elif cells and name in cells: + # 如果value不为空 if value is not None: - raise TypeError(f"The type of value should be cell, but got {type(value).__name__}.") + # 抛出TypeError异常 + raise TypeError(f"For 'Cell', the type of {name} should be cell, but got {type(value).__name__}.") + # 将name和None设置为cells[name] self._cells[name] = None + # 如果value是Tensor类型 elif isinstance(value, Tensor): - if context.get_context("mode") == context.PYNATIVE_MODE: - if name in self.__dict__: - del self.__dict__[name] - tensor_list[name] = value - else: - object.__setattr__(self, name, value) + # 调用_set_attr_for_tensor方法 + self._set_attr_for_tensor(name, value) + # 其他情况 else: + # 如果value是Primitive类型 if isinstance(value, Primitive): + # 调用set_prim_instance_name方法 value.set_prim_instance_name(name) + # 将name和value设置为primitives[name] + self._primitives[name] = value + # 将name和value设置为object.__setattr__(self, name, value) object.__setattr__(self, name, value) + # 如果name不在Cell.IGNORE_LIST中 if name not in Cell.IGNORE_LIST: + # 将_attr_synced设置为False self._attr_synced = False def extend_repr(self): """ - Sets the extended representation of the Cell. + Expand the description of Cell. To print customized extended information, re-implement this method in your own cells. """ @@ -497,6 +1005,9 @@ class Cell(Cell_): return self.__repr__() def __repr__(self): + ''' + 返回一个字符串，该字符串表示当前对象的信息 + ''' extra_str = self.extend_repr() info_str = self.__class__.__name__ + '<' if self._cells: @@ -511,11 +1022,11 @@ class Cell(Cell_): info_str += extra_str + '>' return info_str - def load_parameter_slice(self, params): + def load_parameter_slice(self, params):#利用并行策略将parameters替代为分割的张量 """ Replace parameters with sliced tensors by parallel strategies. - Please refer to the usage in source code of `mindspore.common._Executor.compile`. + Please refer to the usage in source code of `mindspore.common._CellGraphExecutor.compile`. Args: params (dict): The parameters dictionary used for initializing the data graph. @@ -523,19 +1034,28 @@ class Cell(Cell_): if params is None: params = self.parameters_dict() if isinstance(params, OrderedDict): + # 遍历params中的每一个key for key in params: + # 获取key对应的tensor tensor = params[key].data + # 如果key不在parameter_layout_dict中，则输出警告信息 if key not in self.parameter_layout_dict: - logger.info("layout dict does not contain the key %s.", key) + logger.info("The layout dict does not contain the key %s.", key) continue + # 如果key已经被sliced，则输出警告信息 if params[key].sliced: - logger.debug("Param %s is already sliced.", key) + logger.debug("The param %s is already sliced.", key) continue + # 如果key不在parameter_layout_dict中，则输出警告信息 layout = self.parameter_layout_dict[key] + # 使用layout对tensor进行加载 new_tensor = _load_tensor_by_layout(tensor, layout) + # 设置key对应的tensor params[key].set_data(new_tensor, True) else: - raise TypeError("Parameters need OrderedDict type, but got {}.".format(type(params))) + # 如果params的类型不是OrderedDict，则抛出异常 + raise TypeError("For 'load_parameter_slice', the argument 'params' should be OrderedDict type, " + "but got {}.".format(type(params))) def _load_inputs(self, *inputs): """ @@ -548,13 +1068,16 @@ class Cell(Cell_): # judge if *args exists in input if self.argspec[1] is not None: prefix = self.argspec[1] + # 如果inputs中有*args，则将其赋值给prefix for i in range(len(inputs)): key = prefix + str(i) self._construct_inputs_names = self._construct_inputs_names + (key,) self._construct_inputs_num = self._construct_inputs_num + 1 + # 对inputs中的每一个tensor，根据layout进行加载 for i, tensor in enumerate(inputs): key = self._construct_inputs_names[i] # if input is not used, self.parameter_layout_dict may not contain the key + # 如果layout存在，则根据layout加载tensor if key not in self.parameter_layout_dict: logger.warning("Layout dict does not contain the key %s.", key) parallel_inputs_run.append(tensor) @@ -566,7 +1089,7 @@ class Cell(Cell_): def set_parallel_input_with_inputs(self, *inputs): """ - Slice inputs tensors by parallel strategies, and set the sliced inputs to `_parallel_input_run` + Slice inputs tensors by parallel strategies. Args: inputs (tuple): inputs of construct method. @@ -578,88 +1101,209 @@ class Cell(Cell_): from mindspore._extends.parse.parser import get_parse_method_of_class fn = get_parse_method_of_class(self) + # 获取类的parse方法 self.argspec = inspect.getfullargspec(fn) + # 获取类的parse方法的参数信息 self._construct_inputs_num = fn.__code__.co_argcount + # 获取类的parse方法的参数数量 self._construct_inputs_names = fn.__code__.co_varnames - assert self._construct_inputs_num > 0 - assert self._construct_inputs_names[0] == 'self' - assert self._construct_inputs_num - 1 <= len(self._construct_inputs_names) + if self._construct_inputs_num <= 0: + # 如果构造输入的数量小于等于0，抛出异常 + raise ValueError(f"For'set_auto_parallel', the number of inputs must be greater than 0," + f"but got {self._construct_inputs_num}.") + # 如果构造输入的第一个元素不是self，抛出异常 + if self._construct_inputs_names[0]!='self': + raise ValueError(f"First member of fn function must be self, but got {self._construct_inputs_names[0]}") + # 如果构造输入的数量减1大于fn函数成员的数量，抛出异常 + if self._construct_inputs_num - 1 > len(self._construct_inputs_names): + raise ValueError(f"Num of inputs must be greater than num of fn function members, num of inputs is \ + {self._construct_inputs_names - 1}, num of fn function members is {len(self._construct_inputs_names)}") + # 将构造输入的第二个元素从fn函数成员中移除 self._construct_inputs_names = self._construct_inputs_names[1:self._construct_inputs_num] + # 将构造输入的数量减1 self._construct_inputs_num = self._construct_inputs_num - 1 + def set_inputs(self, *inputs): + """ + Save set inputs for computation graph. + + Args: + inputs (tuple): Inputs of the Cell object. + + Examples: + >>> import numpy as np + >>> import mindspore + >>> from mindspore import nn, Tensor, context + >>> + >>> context.set_context(mode=context.GRAPH_MODE, device_target="Ascend") + >>> class reluNet(nn.Cell): + ... def __init__(self): + ... super(reluNet, self).__init__() + ... self.relu = nn.ReLU() + ... def construct(self, x): + ... return self.relu(x) + >>> + >>> net = reluNet() + >>> input_dyn = Tensor(shape=[3, None], dtype=mindspore.float32) + >>> net.set_inputs(input_dyn) + >>> input1 = Tensor(np.random.random([3, 10]), dtype=mindspore.float32) + >>> output = net(input1) + + NOTE: + This is an experimental interface that is subject to change or deletion. + """ + + for ele in self._dynamic_shape_inputs: + if isinstance(ele, (str, int, dict)): + raise TypeError(f"For element in 'set_inputs', the type must be Tensor,\ + but got {type(ele)}.") + self._dynamic_shape_inputs = inputs + + def get_inputs(self): + """ + Returns the dynamic_inputs of a cell object in one network. + + Returns: + inputs (tuple): Inputs of the Cell object. + NOTE: + This is an experimental interface that is subject to change or deletion. + """ + + return self._dynamic_shape_inputs + def compile(self, *inputs): """ - Compiles cell. + Compile Cell as a computation graph, the input must be consistent with the input defined in construct. Args: - inputs (tuple): Input parameters. + inputs (tuple): Inputs of the Cell object. """ - _executor.compile(self, *inputs, phase=self.phase, auto_parallel_mode=self._auto_parallel_mode) + # 如果_dynamic_shape_inputs为None或者_dynamic_shape_inputs[0]为None，则调用_cell_graph_executor.compile函数 + if self._dynamic_shape_inputs is None or self._dynamic_shape_inputs[0] is None: + _cell_graph_executor.compile(self, *inputs, phase=self.phase, auto_parallel_mode=self._auto_parallel_mode) + # 如果_dynamic_shape_inputs不为None，但是_dynamic_shape_inputs[0]不为None，则检查_dynamic_shape_inputs + else: + self._check_compile_dynamic_shape(*inputs) + # 如果saved_dynamic_shape不为空，且saved_dynamic_shape[i].shape不等于_dynamic_shape_inputs[i].shape和_dynamic_shape_inputs[i].shape，则终止 + if self.saved_dynamic_shape: + for i in range(len(self.saved_dynamic_shape)): + if self.saved_dynamic_shape[i].shape!= self._dynamic_shape_inputs[i].shape \ + and self.saved_dynamic_shape[i].shape!= self._dynamic_shape_inputs[i].shape: + break + return + # 否则，将_dynamic_shape_inputs赋值给saved_dynamic_shape + self.saved_dynamic_shape = self._dynamic_shape_inputs + # 调用_cell_graph_executor.compile函数，并设置phase为self.phase，auto_parallel_mode为self._auto_parallel_mode + _cell_graph_executor.compile(self, *self._dynamic_shape_inputs, phase=self.phase, + auto_parallel_mode=self._auto_parallel_mode) + # 打印日志 + logger.debug("Compiled Graph with dynamic shape") def compile_and_run(self, *inputs): """ - Compiles and runs cell. + Compile and run Cell, the input must be consistent with the input defined in construct. + + Note: It is not recommended to call directly. Args: - inputs (tuple): Input parameters. + inputs (tuple): Inputs of the Cell object. Returns: Object, the result of executing. """ + # 设置自动并行编译和运行 self._auto_parallel_compile_and_run = True + # 编译和运行 self.compile(*inputs) + # 创建新的输入 new_inputs = [] for i in inputs: + # 如果输入是Tensor类型 if isinstance(i, Tensor): + # 如果输入有初始化 + if i.has_init: + # 初始化输入 + i.init_data() + # 将输入添加到新的输入列表中 new_inputs.append(i) + # 如果输入是COOTensor或CSRTensor类型 + elif isinstance(i, (COOTensor, CSRTensor)): + # 将输入添加到新的输入列表中 + new_inputs.append(i) + # 如果输入是Variable类型 + elif isinstance(i, Variable): + # 将输入添加到新的输入列表中 + new_inputs.append(i.value) + # 如果输入是context.get_context("grad_for_scalar")为True的时候且输入是int或float类型 elif context.get_context("grad_for_scalar") and isinstance(i, (int, float)): + # 将输入添加到新的输入列表中 + new_inputs.append(i) + # 如果自动并行模式为True + elif hasattr(self, "enable_tuple_broaden") and self.enable_tuple_broaden and isinstance(i, tuple) and \ + _check_all_tensor(i): + # 将输入添加到新的输入列表中 new_inputs.append(i) + # 如果自动并行模式为True if self._auto_parallel_mode: + # 如果新的输入是Tensor类型，且输入的虚拟标志为True if new_inputs and isinstance(new_inputs[0], Tensor) and inputs[0].virtual_flag: - # get parallel inputs in sink mode, parallel inputs set in _executor.compile + # get parallel inputs in sink mode, parallel inputs set in _cell_graph_executor.compile + # 获取并行输入 parallel_inputs_run = self._parallel_inputs_run + # 否则 else: + # 将新的输入添加到并行输入列表中 parallel_inputs_run = new_inputs - return _executor(self, *parallel_inputs_run, phase=self.phase) - return _executor(self, *new_inputs, phase=self.phase) + # 运行_cell_graph_executor + return _cell_graph_executor(self, *parallel_inputs_run, phase=self.phase) + # 否则 + return _cell_graph_executor(self, *new_inputs, phase=self.phase) def auto_parallel_compile_and_run(self): + """ + Whether or not to execute compile and run in 'AUTO_PARALLEL' or 'SEMI_AUTO_PARALLEL' mode. + + Returns: + bool, `_auto_parallel_compile_and_run` value. + """ return self._auto_parallel_compile_and_run def exec_checkpoint_graph(self): """Executes saving checkpoint graph operation.""" - _executor(self, phase='save') + # 调用_cell_graph_executor函数，传入参数phase='save' + _cell_graph_executor(self, phase='save') - def insert_param_to_cell(self, param_name, param, check_name=True): + def insert_param_to_cell(self, param_name, param, check_name_contain_dot=True):#给当前cell添加parameter """ Adds a parameter to the current cell. - Inserts a parameter with given name to the cell. Please refer to the usage in - source code of `mindspore.nn.Cell.__setattr__`. + Inserts a parameter with given name to the cell. The method is currently used in + `mindspore.nn.Cell.__setattr__`. Args: param_name (str): Name of the parameter. param (Parameter): Parameter to be inserted to the cell. - check_name (bool): Determines whether the name input is compatible. Default: True. + check_name_contain_dot (bool): Determines whether the name input is compatible. Default: True. Raises: KeyError: If the name of parameter is null or contains dot. - AttributeError: If user did not call init() first. TypeError: If the type of parameter is not Parameter. """ if not param_name: - raise KeyError("The name of parameter should not be null.") - if check_name and '.' in param_name: - raise KeyError("The name of parameter should not contain \".\"") + raise KeyError("For 'insert_param_to_cell', the argument 'param_name' should not be None.") + if check_name_contain_dot and '.' in param_name: + raise KeyError("For 'insert_param_to_cell', the argument 'param_name' should not contain \".\"") if '_params' not in self.__dict__: - raise AttributeError("You need call init() first.") + raise AttributeError("For 'insert_param_to_cell', please call Cell.__init__() firstly.") if hasattr(self, param_name) and param_name not in self._params: - raise KeyError("Duplicated parameter name '{}'.".format(param_name)) + raise KeyError("For 'insert_param_to_cell', the {} parameter already exists in the network. Cannot " + "insert another parameter with the same name.".format(param_name)) if not isinstance(param, Parameter) and param is not None: - raise TypeError("The type of parameter should be 'Parameter' if not None.") + raise TypeError(f"For 'insert_param_to_cell', the argument 'param' should be 'Parameter' if not None, " + f"but got {type(param)}.") if isinstance(param, Parameter) and param.name == PARAMETER_NAME_DEFAULT: param.name = param_name self._params[param_name] = param @@ -668,20 +1312,28 @@ class Cell(Cell_): """ Cast parameter according to auto mix precision level in pynative mode. + This interface is currently used in the case of auto mix precision and usually needs not to be used explicitly. + Args: - param (Parameter): The parameter to cast. + param (Parameter): Parameters, the type of which should be cast. + + Returns: + Parameter, the input parameter with type automatically cast. """ - if hasattr(self, "_mindspore_flags"): - if self._mindspore_flags.get('fp32'): + mixed_type = self.get_mixed_precision_type() + # 如果mixed_type不为NOTSET，则设置param的cast_dtype + if mixed_type!= MixedPrecisionType.NOTSET: + if mixed_type == MixedPrecisionType.FP32: param.set_cast_dtype(mstype.float32) - elif self._mindspore_flags.get('fp16'): + elif mixed_type == MixedPrecisionType.FP16: param.set_cast_dtype(mstype.float16) - elif hasattr(param, "set_cast_dtype"): - # retest dtype - param.set_cast_dtype() + # 如果param有set_cast_dtype方法，则重新设置param的cast_dtype + elif hasattr(param, "set_cast_dtype"): + # retest dtype + param.set_cast_dtype() return param - def insert_child_to_cell(self, child_name, child_cell): + def insert_child_to_cell(self, child_name, child_cell):#用特定名称给当前cell添加子cell """ Adds a child cell to the current cell with a given name. @@ -694,88 +1346,153 @@ class Cell(Cell_): TypeError: Child Cell's type is incorrect. """ if not child_name or '.' in child_name: - raise KeyError("Child cell name is incorrect.") + raise KeyError("For 'insert_child_to_cell', the argument 'child_name' should not be None and " + "should not contain '.'") if hasattr(self, child_name) and child_name not in self._cells: - raise KeyError("Duplicate child name '{}'.".format(child_name)) + raise KeyError("For 'insert_child_to_cell', the {} child cell already exists in the network. Cannot " + "insert another child cell with the same name.".format(child_name)) if not isinstance(child_cell, Cell) and child_cell is not None: - raise TypeError("Child cell type is incorrect.") + raise TypeError(f"For 'insert_child_to_cell', the argument 'child_cell' should be 'Cell' if not None, " + f"but got type {type(child_cell)}.") self._cells[child_name] = child_cell def construct(self, *inputs, **kwargs): """ Defines the computation to be performed. This method must be overridden by all subclasses. + Note: It is not supported currently that inputs contain both tuple and non-tuple types at same time. + + Args: + inputs (tuple): Tuple of variable parameters. + kwargs (dict): Dictionary of variable keyword parameters. + Returns: Tensor, returns the computed result. """ return None - def remove_redundant_parameters(self): - """Remove the redundant parameters""" + def remove_redundant_parameters(self):#移除多余的parameters + """ + Remove the redundant parameters. + + This interface usually needs not to be used explicitly. + """ cells = self.cells_and_names() + # 遍历cells中的每一个cell for _, cell in cells: + # 获取cell的参数 params = cell._params.items() + # 遍历cell的参数 for param_name, param in list(params): + # 如果参数名在parallel_parameter_name_list中，则从参数中移除 if param.name not in self.parallel_parameter_name_list: cell._params.pop(param_name) logger.info("remove the redundant parameter: %s", param.name) continue + # 获取cell的字典 cell_dict = cell.__dict__ + # 遍历cell的字典 for key in cell_dict: + # 如果字典中的值是ParameterTuple类型，则获取其中的参数 if isinstance(cell_dict[key], ParameterTuple): param_tuple = cell_dict[key] + # 创建新的参数列表 new_param_tuple = [] + # 遍历参数列表 for param in param_tuple: + # 如果参数名在parallel_parameter_name_list中，则从参数列表中移除 if param.name not in self.parallel_parameter_name_list: logger.info("remove the redundant parameter: %s in ParameterTuple", param.name) continue + # 将参数添加到新的参数列表中 new_param_tuple.append(param) + # 将新的参数列表赋值给cell的字典中的值 cell.__dict__[key] = ParameterTuple(new_param_tuple) - def init_parameters_data(self, auto_parallel_mode=False): + def init_parameters_data(self, auto_parallel_mode=False):#初始化_parameters_data """ Initialize all parameters and replace the original saved parameters in cell. Note: trainable_params() and other similar interfaces may return different parameter instance after - `init_parameters_data`, do not save these result. + `init_parameters_data`, do not save these results. Args: - auto_parallel_mode (bool): If running in auto_parallel_mode. + auto_parallel_mode (bool): If running in auto_parallel_mode. Default: False. Returns: Dict[Parameter, Parameter], returns a dict of original parameter and replaced parameter. """ replace = dict() + # 定义一个函数_updata，用于更新参数 def _updata(param): + # 如果param在replace中，则返回replace中的值 if param in replace: - return replace[param] + return replace.get(param) + # 定义一个变量layout，用于存储更新后的参数 layout = None + # 定义一个变量set_sliced，用于标记是否设置sliced set_sliced = False + # 如果auto_parallel_mode为True，则设置set_sliced为True if auto_parallel_mode: set_sliced = True + # 如果param的name不在self.parameter_layout_dict中，则输出警告信息 if param.name not in self.parameter_layout_dict: logger.debug("Layout dict does not contain the key %s.", param.name) + # 否则，获取layout else: layout = self.parameter_layout_dict[param.name] + # 将param的初始数据更新到layout中 new_p = param.init_data(layout, set_sliced=set_sliced) + # 将param更新到replace中 replace[param] = new_p + # 返回更新后的参数 return new_p # replace all original usage. + # 更新所有原有的usage cells = self.cells_and_names() for _, cell in cells: + # 获取cell的参数 params = cell._params.items() + # 遍历参数 for param_name, param in params: + # 如果auto_parallel_mode为False，则跳过 + if not auto_parallel_mode: + # 将参数更新到cell的参数中 + cell._params[param_name] = _updata(param) + continue + # 如果param的name在parallel_parameter_name_list中，则跳过 + if param.name in self.parallel_parameter_name_list: + # 将参数更新到cell的参数中 + cell._params[param_name] = _updata(param) + continue + # 否则，将参数更新到cell的参数中 cell._params[param_name] = _updata(param) + # 获取cell的__dict__ cell_dict = cell.__dict__ + # 遍历cell的__dict__ for key in cell_dict: + # 如果cell的__dict__中是ParameterTuple类型 if isinstance(cell_dict[key], ParameterTuple): + # 将ParameterTuple类型的参数更新到cell的__dict__中 param_tuple = cell_dict[key] new_param_tuple = [] + # 遍历ParameterTuple的参数 for param in param_tuple: - new_param_tuple.append(_updata(param)) + # 如果auto_parallel_mode为False，则跳过 + if not auto_parallel_mode: + # 将参数更新到new_param_tuple中 + new_param_tuple.append(_updata(param)) + continue + # 如果param的name在parallel_parameter_name_list中，则跳过 + if param.name in self.parallel_parameter_name_list: + # 将参数更新到new_param_tuple中 + new_param_tuple.append(_updata(param)) + # 否则，将参数更新到new_param_tuple中 + else: + new_param_tuple.append(param) cell.__dict__[key] = ParameterTuple(new_param_tuple) return replace @@ -792,17 +1509,33 @@ class Cell(Cell_): OrderedDict, return parameters dictionary. """ param_dict = OrderedDict() + # 遍历获取的参数 for param in self.get_parameters(expand=recurse): + # 将参数名和参数值存入字典中 param_dict[param.name] = param + # 返回字典 return param_dict def parameters_broadcast_dict(self, recurse=True): + """ + Gets the parameters broadcast dictionary of this cell. + + Args: + recurse (bool): Whether contains the parameters of subcells. Default: True. + + Returns: + OrderedDict, return parameters broadcast dictionary. + """ param_dict = OrderedDict() + # 遍历获取的参数 for param in self.get_parameters(expand=recurse): + # 如果参数不是层级并行，则将其存入字典中 if param.layerwise_parallel is False: param_dict[param.name] = param + # 如果字典为空，则返回None if not param_dict: return None + # 否则返回字典 return param_dict def update_parameters_name(self, prefix='', recurse=True): @@ -812,17 +1545,43 @@ class Cell(Cell_): Adds the given prefix to the names of parameters. Args: - prefix (str): The prefix string. + prefix (str): The prefix string. Default: ''. recurse (bool): Whether contains the parameters of subcells. Default: True. """ Validator.check_str_by_regular(prefix) for name, param in self.parameters_and_names(expand=recurse): - if prefix != '': + # 如果prefix不为空，则将param的is_init设置为False + if prefix!= '': param.is_init = False + # 将prefix加上param的name赋值给param的name param.name = prefix + name - def trainable_params(self, recurse=True): + def _update_local_parameters_name(self, prefix='', recurse=True): + """ + Updates the names of local parameters with given prefix string. + + Adds the given prefix to the names of local parameters. + + Local parameters means the parameters without user input. + + Args: + prefix (str): The prefix string. Default: ''. + recurse (bool): Whether contains the parameters of subcells. Default: True. + """ + + Validator.check_str_by_regular(prefix) + for name, param in self.parameters_and_names(expand=recurse): + # 如果name在_user_parameters中，则跳过 + if name in self._user_parameters: + continue + # 如果prefix不为空，则将param设置为is_init为False + if prefix!= '': + param.is_init = False + # 将prefix加上name赋值给param.name + param.name = prefix + name + + def trainable_params(self, recurse=True):#返回全部的trainable_params """ Returns all trainable parameters. @@ -836,7 +1595,7 @@ class Cell(Cell_): """ return list(filter(lambda x: x.requires_grad, self.get_parameters(expand=recurse))) - def untrainable_params(self, recurse=True): + def untrainable_params(self, recurse=True):#返回全部的untrainable_params """ Returns all untrainable parameters. @@ -850,18 +1609,22 @@ class Cell(Cell_): """ return list(filter(lambda x: not x.requires_grad, self.get_parameters(expand=recurse))) - def get_parameters(self, expand=True): + def get_parameters(self, expand=True):#返回cell上的迭代器 """ Returns an iterator over cell parameters. - Yields parameters of this cell. If `expand` is True, yield parameters of this cell and all subcells. + Yields parameters of this cell. If `expand` is true, yield parameters of this cell and all subcells. Args: expand (bool): If true, yields parameters of this cell and all subcells. Otherwise, only yield parameters that are direct members of this cell. Default: True. + Returns: + Iteration, all parameters at the cell. + Examples: - >>> net = Net() + >>> from mindspore import nn + >>> net = nn.Dense(3, 4) >>> parameters = [] >>> for item in net.get_parameters(): ... parameters.append(item) @@ -870,52 +1633,70 @@ class Cell(Cell_): yield param def check_names(self): + """ + Check the names of cell parameters. + """ names = set("") + # 创建一个空集合，用于存储参数名称 for value, param in self.parameters_and_names(): + # 遍历参数和参数名称 if param.name in names: - raise ValueError("The value of {} is {}, its name '{}' already exists.". - format(value, param, param.name)) + # 如果参数名称已存在，抛出异常 + raise ValueError("The value of {} is {}, its name '{}' already exists. " + "Please set a unique name for the parameter.".format(value, param, param.name)) names.add(param.name) - def parameters_and_names(self, name_prefix='', expand=True): + def parameters_and_names(self, name_prefix='', expand=True):#返回cell上的迭代器包括名称 """ Returns an iterator over cell parameters. - Includes the parameter's name and itself. + Includes the parameter's name and itself. Args: name_prefix (str): Namespace. Default: ''. expand (bool): If true, yields parameters of this cell and all subcells. Otherwise, only yield parameters that are direct members of this cell. Default: True. + Returns: + Iteration, all the names and corresponding parameters in the cell. + Examples: - >>> n = Net() + >>> from mindspore import nn + >>> n = nn.Dense(3, 4) >>> names = [] >>> for m in n.parameters_and_names(): ... if m[0]: ... names.append(m[0]) """ cells = [] + # 如果expand为True，则cells为self.cells_and_names(name_prefix=name_prefix) if expand: cells = self.cells_and_names(name_prefix=name_prefix) + # 否则cells为[(name_prefix, self)] else: cells.append((name_prefix, self)) params_set = set() + # 遍历cells，获取每一个cell的参数 for cell_name, cell in cells: params = cell._params.items() + # 遍历参数 for par_name, par in params: + # 如果参数未初始化，则获取参数 if par.inited_param is not None: par = par.inited_param + # 如果参数不为None，且id不在params_set中，则将参数添加到params_set中 if par is not None and id(par) not in params_set: params_set.add(id(par)) par_new_name = par_name + # 如果cell_name为空，则将参数的名称添加到par_new_name中 if cell_name: par_new_name = cell_name + '.' + par_new_name + # 返回参数的新名称和参数 yield par_new_name, par - def cells_and_names(self, cells=None, name_prefix=''): + def cells_and_names(self, cells=None, name_prefix=''):#返回网络中cell上的迭代器和名字 """ Returns an iterator over all cells in the network. @@ -925,9 +1706,20 @@ class Cell(Cell_): cells (str): Cells to iterate over. Default: None. name_prefix (str): Namespace. Default: ''. + Returns: + Iteration, all the child cells and corresponding names in the cell. + Examples: - >>> n = Net() + >>> from mindspore import nn + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.conv = nn.Conv2d(3, 64, 3) + ... def construct(self, x): + ... out = self.conv(x) + ... return out >>> names = [] + >>> n = Net() >>> for m in n.cells_and_names(): ... if m[0]: ... names.append(m[0]) @@ -947,18 +1739,23 @@ class Cell(Cell_): for ele in cell.cells_and_names(t_cells, cells_name_prefix): yield ele - def cells(self): - """Returns an iterator over immediate cells.""" + def cells(self):#自定义 + """ + Returns an iterator over immediate cells. + + Returns: + Iteration, the immediate cells in the cell. + """ return self.name_cells().values() - def _set_scope(self, name): + def _set_scope(self, name):#首次设置名称 """Sets the name on the first time.""" if self._scope is None: self._scope = name elif self._scope == 'recompute_': self._scope = self._scope + name - def _children_scope_recursive(self, parent_prefix='Default'): + def _children_scope_recursive(self, parent_prefix='Default'):#递归地生成网络的每一层 """Generates the scope of each layer of the network recursively.""" reserve_class_name_in_scope = context.get_context("reserve_class_name_in_scope") @@ -972,105 +1769,204 @@ class Cell(Cell_): if reserve_class_name_in_scope else "")): yield key, value - def get_scope(self): - """Returns the scope of a cell object in one network.""" + def get_scope(self):#返回一个网络中的cell对象的scope + """ + Returns the scope of a cell object in one network. + + Returns: + String, scope of the cell. + """ return self._scope - def generate_scope(self): + def generate_scope(self):#为网络中的每个cell对象生成scope """Generate the scope for each cell object in the network.""" for name, cell in self._children_scope_recursive(): cell._set_scope(name) def name_cells(self): """ - Returns an iterator over all cells in the network. + Returns an iterator over all immediate cells in the network. Include name of the cell and cell itself. + + Returns: + Dict, all the child cells and corresponding names in the cell. """ value_set = set() cells = OrderedDict() + # 遍历cells字典，获取每一个cell的值 for name, cell in self._cells.items(): + # 如果cell不为空，且不在value_set中，则将cell添加到value_set中 if cell is not None and cell not in value_set: value_set.add(cell) cells[name] = cell + # 返回cells字典 return cells + def _add_mixed_precision_flag(self, **flags): + """Add mixed precision flag to current cell""" + # 如果flags中存在fp16，且flags中的fp16属性为True + if "fp16" in flags and flags.get("fp16", False): + # 设置当前cell的mixed precision类型为FP16 + Cell_.set_mixed_precision_type(self, MixedPrecisionType.FP16) + # 如果flags中存在fp32，且flags中的fp32属性为True + if "fp32" in flags and flags.get("fp32", False): + # 设置当前cell的mixed precision类型为FP32 + Cell_.set_mixed_precision_type(self, MixedPrecisionType.FP32) + + def _add_mixed_precision_flag_recursive(self, **flags): + """Add mixed precision flag to each cell""" + # 如果flags中存在fp16，且flags中的fp16属性为True + if "fp16" in flags and flags.get("fp16", False): + # 调用_set_mixed_precision_type_recursive函数，设置当前cell的mixed precision类型为FP16 + self._set_mixed_precision_type_recursive(MixedPrecisionType.FP16) + # 如果flags中存在fp32，且flags中的fp32属性为True + if "fp32" in flags and flags.get("fp32", False): + # 调用_set_mixed_precision_type_recursive函数，设置当前cell的mixed precision类型为FP32 + self._set_mixed_precision_type_recursive(MixedPrecisionType.FP32) + def add_flags(self, **flags): + """ + Add customized attributes for cell. + + This method is also called when the cell class is instantiated and the class parameter 'flags' is set to True. + + Args: + flags (dict): Network configuration information, currently it is used for the binding of network and + dataset. Users can also customize network attributes by this parameter. Default: None. + """ if not hasattr(self, "_mindspore_flags"): + # 如果_mindspore_flags不存在，则初始化为空字典 self._mindspore_flags = {} + # 将flags添加到_mindspore_flags中 self._mindspore_flags.update({**flags}) + # 将flags的属性添加到self中 self.__dict__.update({**flags}) + # 将flags的属性添加到self中，并将其设置为self + self._add_mixed_precision_flag(**flags) + # 返回self return self def add_flags_recursive(self, **flags): + """ + If a cell contains child cells, this method can recursively customize attributes of all cells. + + Args: + flags (dict): Network configuration information, currently it is used for the binding of network and + dataset. Users can also customize network attributes by this parameter. Default: None. + """ self.add_flags(**flags) - if hasattr(self, '_cell_init_args'): - self._cell_init_args += str({**flags}) + # 添加flags + self._add_mixed_precision_flag_recursive(**flags) + # 递归添加mixed precision flag for cell in self.cells(): + # 遍历cells cell.add_flags_recursive(**flags) + # 递归添加flags return self + def _add_init_args(self, **args): + if hasattr(self, '_cell_init_args'): + self._cell_init_args += str({**args}) + def get_flags(self): + """ + Get the self_defined attributes of the cell, which can be added by `add_flags` method. + """ if not hasattr(self, "_mindspore_flags"): self._mindspore_flags = {} return self._mindspore_flags + def _set_mixed_precision_type_recursive(self, mixed_type): + """Set mixed precision type to each cell""" + Cell_.set_mixed_precision_type(self, mixed_type) + # 遍历每一个Cell + for cell in self.cells(): + # 递归调用_set_mixed_precision_type_recursive函数 + cell._set_mixed_precision_type_recursive(mixed_type) + def to_float(self, dst_type): """ Add cast on all inputs of cell and child cells to run with certain float type. - If `dst_type is mindspore.dtype.float16`, all the inputs of Cell including input, Parameter, Tensor - as const will be cast to float16. Please refer to the usage in source code of - `mindspore.train.amp.build_train_network`. + If `dst_type` is `mindspore.dtype.float16`, all the inputs of Cell, including input, Parameter and Tensor, will + be cast to float16. Please refer to the usage in source code of :func:`mindspore.build_train_network`. Note: Multiple calls will overwrite. Args: - dst_type (:class:`mindspore.dtype`): Transfer Cell to Run with dst_type. - dst_type can be `mindspore.dtype.float16` or `mindspore.dtype.float32`. + dst_type (:class:`mindspore.dtype`): Transfer cell to run with dst_type. + dst_type can be `mstype.float16` or `mstype.float32`. + + Returns: + Cell, the cell itself. Raises: - ValueError: If dst_type is not float32 nor float16. + ValueError: If dst_type is not mstype.float32 or mstype.float16. + + Examples: + >>> import mindspore.nn as nn + >>> from mindspore import dtype as mstype + >>> + >>> net = nn.Conv2d(120, 240, 4, has_bias=False, weight_init='normal') + >>> net.to_float(mstype.float16) """ if dst_type not in (mstype.float16, mstype.float32): - raise ValueError("dst_type should inside float32 or float16.") + raise ValueError("For 'to_float', the argument 'dst_type' should be float32 or float16, " + "but got {}.".format(dst_type)) + if dst_type == mstype.float16: + # 设置为FP16类型 + self._set_mixed_precision_type_recursive(MixedPrecisionType.FP16) + else: + # 设置为FP32类型 + self._set_mixed_precision_type_recursive(MixedPrecisionType.FP32) + # 将flags字典中的值设置为dst_type的类型 flags = {'fp16': dst_type == mstype.float16, 'fp32': dst_type == mstype.float32} - self.add_flags_recursive(**flags) + # 添加初始化参数 + self._add_init_args(**flags) return self - def set_acc(self, acc_type): + def set_boost(self, boost_type): """ In order to improve the network performance, configure the network auto enable to accelerate the algorithm in the algorithm library. - If `acc_type is not in the algorithm library`, Please view the algorithm in the algorithm library - through `algorithm library`. + If `boost_type` is not in the algorithm library. Please view the algorithm in the algorithm library through + `algorithm library `_. Note: Some acceleration algorithms may affect the accuracy of the network, please choose carefully. Args: - acc_type (str): accelerate algorithm. + boost_type (str): accelerate algorithm. + + Returns: + Cell, the cell itself. Raises: - ValueError: If acc_type is not in the algorithm library. + ValueError: If boost_type is not in the algorithm library. """ - if acc_type not in ("less_bn",): - raise ValueError("acc_type is not in the algorithm library.") - flags = {"less_bn": acc_type == "less_bn"} + if boost_type not in ("less_bn",): + raise ValueError("For 'set_boost', the argument 'boost_type' should be 'less_bn', " + "but got {}.".format(boost_type)) + flags = {"less_bn": boost_type == "less_bn"} + # 调用add_flags_recursive函数，传入参数flags self.add_flags_recursive(**flags) return self def set_grad(self, requires_grad=True): """ - Sets the cell flag for gradient. In pynative mode, this parameter specifies whether the network require - gradients. If True, the backward network needed to compute the gradients will be generated when the forward + Sets the cell flag for gradient. In pynative mode, this parameter specifies whether the network requires + gradients. If true, the backward network needed to compute the gradients will be generated when the forward network is executed. Args: requires_grad (bool): Specifies if the net need to grad, if it is - True, cell will construct backward network in pynative mode. Default: True. + true, the cell will construct backward network in pynative mode. Default: True. + + Returns: + Cell, the cell itself. """ self.requires_grad = requires_grad return self @@ -1081,31 +1977,35 @@ class Cell(Cell_): The cell itself and all children cells will be set to training mode. Layers that have different constructions for training and predicting, such as `BatchNorm`, will distinguish between the branches by this attribute. If - set to True, the training branch will be executed, otherwise another branch. + set to true, the training branch will be executed, otherwise another branch. Args: mode (bool): Specifies whether the model is training. Default: True. + + Returns: + Cell, the cell itself. """ if mode is False: self._phase = 'predict' else: self._phase = 'train' + # 将cell的training属性设置为mode self.add_flags_recursive(training=mode) return self def set_broadcast_flag(self, mode=True): """ - Set the cell to data_parallel mode. - - The cell can be accessed as an attribute using the given name. + Set parameter broadcast mode for this cell. Args: - mode (bool): Specifies whether the model is data_parallel. Default: True. + mode (bool): Specifies whether the mode is parameter broadcast. Default: True. """ + # 调用add_flags_recursive函数，传入参数broadcast_flag，并将mode设置为True self.add_flags_recursive(broadcast_flag=mode) + # 返回self return self - def set_auto_parallel(self): + def set_auto_parallel(self):#设置自动并行 """ Set the cell to auto parallel mode. @@ -1117,31 +2017,348 @@ class Cell(Cell_): self.add_flags(auto_parallel=True) self._get_construct_inputs_number_and_name() - def _hook_construct(self, *inputs, **kwargs): - """Hook construct method to replace original construct method when hook function enabled.""" - inputs = self._backward_hook(*inputs) - inputs = self.construct(inputs) - outputs = self._backward_hook(inputs) - return outputs + def flatten_weights(self): + """ + Reset data for weight parameters so that they are using contiguous memory chunks grouped by data type. + """ + Tensor._flatten_tensors(self.trainable_params()) # pylint: disable=W0212 - def register_backward_hook(self, fn): + def _run_forward_pre_hook(self, inputs): """ - Set the cell backward hook function. Note that this function is only supported in Pynative Mode. + Running forward pre hook function registered on Cell object. + + Args: + inputs: The input objects of cell object. + + Returns: + - **outputs** - New input objects or none. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + """ + # 获取Cell对象的类名 + cell_id = self.cls_name + "(" + str(id(self)) + ")" + # 遍历forward_pre_hook字典 + for fn in self._forward_pre_hook.values(): + # 调用fn函数 + ret = fn(cell_id, inputs) + # 如果返回值不为空，则将其转换为tuple + if ret is not None: + if not isinstance(ret, tuple): + inputs = (ret,) + else: + inputs = ret + # 返回新的输入对象 + return inputs + + def register_forward_pre_hook(self, hook_fn): + """ + Register forward pre hook function for Cell object. Note: - fn must be defined as the following code. `cell_name` is the name of registered cell. - `grad_input` is gradient passed to the cell. `grad_output` is the gradient computed and passed to the - next cell or primitive, which may be modified and returned. - hook_fn(cell_name, grad_input, grad_output) -> Tensor or None. + - The `register_forward_pre_hook(hook_fn)` does not work in graph mode or ms_function. + - 'hook_fn' must be defined as the following code. + `cell_id` is the information of registered Cell object, including name and ID. `inputs` is the forward + input objects passed to the Cell. The 'hook_fn' can modify the forward input objects by returning new + forward input objects. + - It should have the following signature: + hook_fn(cell_id, inputs) -> new input objects or none. + - In order to prevent running failed when switching to graph mode, it is not recommended to write it in the + `construct` function of Cell object. In the pynative mode, if the `register_forward_pre_hook` function is + called in the `construct` function of the Cell object, a hook function will be added at each run time of + Cell object. + + Args: + hook_fn (function): Python function. Forward pre hook function. + + Returns: + Handle, it is an instance of `mindspore.common.hook_handle.HookHandle` and corresponding to the `hook_fn` . + The handle can be used to remove the added `hook_fn` by calling `handle.remove()` . + + Raises: + TypeError: If the `hook_fn` is not a function of python. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + + Examples: + >>> import numpy as np + >>> import mindspore + >>> import mindspore.nn as nn + >>> from mindspore import Tensor + >>> from mindspore import context + >>> from mindspore.ops import GradOperation + >>> context.set_context(mode=context.PYNATIVE_MODE) + >>> def forward_pre_hook_fn(cell_id, inputs): + ... print("forward inputs: ", inputs) + ... + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.mul = nn.MatMul() + ... self.handle = self.mul.register_forward_pre_hook(forward_pre_hook_fn) + ... + ... def construct(self, x, y): + ... x = x + x + ... x = self.mul(x, y) + ... return x + >>> grad = GradOperation(get_all=True) + >>> net = Net() + >>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32))) + forward inputs: (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], + dtype=Float32, value= [ 1.00000000e+00])) + >>> print(output) + (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32, + value= [ 2.00000000e+00])) + """ + if context.get_context("mode") != context.PYNATIVE_MODE: + logger.warning(f"'register_forward_pre_hook' function is only supported in pynative mode, you can use " + f"context.set_context to set pynative mode.") + return HookHandle() + + if not isinstance(hook_fn, (FunctionType, MethodType)): + raise TypeError(f"When using'register_forward_pre_hook(hook_fn)', the type of 'hook_fn' should be python " + f"function, but got {type(hook_fn)}.") + if hook_fn.__code__.co_name == "staging_specialize": + raise TypeError(f"Decorating hook function {hook_fn.__name__} with '@ms_function' is not supported.") + + # 将hook_fn的名称设置为enable_forward_pre_hook + self._enable_forward_pre_hook = True + # 设置hook_fn的变更 + _pynative_executor.set_hook_changed(self) + if not hasattr(self, '_forward_pre_hook_key'): + # 如果没有_forward_pre_hook_key属性，则设置_forward_pre_hook_key为-1 + self._forward_pre_hook_key = -1 + # 将_forward_pre_hook_key加1 + self._forward_pre_hook_key += 1 + # 将hook_fn添加到_forward_pre_hook中 + self._forward_pre_hook[self._forward_pre_hook_key] = hook_fn + # 返回HookHandle + handle = HookHandle(self, self._forward_pre_hook_key, "_forward_pre_hook") + return handle + + def _run_forward_hook(self, inputs, output): + """ + Running forward hook function registered on Cell object. + + Args: + inputs: The input objects of Cell object. + output: The output object of Cell object. + + Returns: + - **output** - New output object or none. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + """ + cell_id = self.cls_name + "(" + str(id(self)) + ")" + # 遍历_forward_hook字典，调用fn函数，返回输出 + for fn in self._forward_hook.values(): + ret = fn(cell_id, inputs, output) + if ret is not None: + output = ret + return output + + def register_forward_hook(self, hook_fn): + """ + Set the Cell forward hook function. + + Note: + - The `register_forward_hook(hook_fn)` does not work in graph mode or ms_function. + - 'hook_fn' must be defined as the following code. + `cell_id` is the information of registered Cell object, including name and ID. `inputs` is the forward + input objects passed to the Cell. `output` is the forward output object of the Cell. The 'hook_fn' can + modify the forward output object by returning new forward output object. + - It should have the following signature: + hook_fn(cell_id, inputs, output) -> new output object or none. + - In order to prevent running failed when switching to graph mode, it is not recommended to write it in the + `construct` function of Cell object. In the pynative mode, if the `register_forward_hook` function is + called in the `construct` function of the Cell object, a hook function will be added at each run time of + Cell object. + + Args: + hook_fn (function): Python function. Forward hook function. + + Returns: + Handle, it is an instance of `mindspore.common.hook_handle.HookHandle` and corresponding to the `hook_fn` . + The handle can be used to remove the added `hook_fn` by calling `handle.remove()` . + + Raises: + TypeError: If the `hook_fn` is not a function of python. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + + Examples: + >>> import numpy as np + >>> import mindspore + >>> import mindspore.nn as nn + >>> from mindspore import Tensor + >>> from mindspore import context + >>> from mindspore.ops import GradOperation + >>> context.set_context(mode=context.PYNATIVE_MODE) + >>> def forward_hook_fn(cell_id, inputs, output): + ... print("forward inputs: ", inputs) + ... print("forward output: ", output) + ... + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.mul = nn.MatMul() + ... self.handle = self.mul.register_forward_hook(forward_hook_fn) + ... + ... def construct(self, x, y): + ... x = x + x + ... x = self.mul(x, y) + ... return x + >>> grad = GradOperation(get_all=True) + >>> net = Net() + >>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32))) + forward inputs: (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], + dtype=Float32, value= [ 1.00000000e+00])) + forward output: 2.0 + >>> print(output) + (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32, + value= [ 2.00000000e+00])) + """ + if context.get_context("mode") != context.PYNATIVE_MODE: + logger.warning(f"'register_forward_hook' function is only supported in pynative mode, you can use " + f"context.set_context to set pynative mode.") + return HookHandle() + + if not isinstance(hook_fn, (FunctionType, MethodType)): + raise TypeError(f"When using'register_forward_hook(hook_fn)', the type of 'hook_fn' should be python " + f"function, but got {type(hook_fn)}.") + if hook_fn.__code__.co_name == "staging_specialize": + raise TypeError(f"Decorating hook function {hook_fn.__name__} with '@ms_function' is not supported.") + + # 注册前向调用的hook函数 + self._enable_forward_hook = True + _pynative_executor.set_hook_changed(self) + if not hasattr(self, '_forward_hook_key'): + self._forward_hook_key = -1 + self._forward_hook_key += 1 + self._forward_hook[self._forward_hook_key] = hook_fn + handle = HookHandle(self, self._forward_hook_key, "_forward_hook") + return handle + + def _backward_hook_construct(self, *inputs): + """ + Backward hook construct method to replace original construct method. Args: - fn (function): Specifies the hook function with grad as input. + inputs: The input objects of Cell object. + + Returns: + - **outputs** - The output objects of Cell object. + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` """ - self._backward_hook = HookBackward(fn, self.cls_name + "(" + str(id(self)) + ")") - self.enable_hook = True + # 如果输入的数量大于1，则将输入转换为Cell的输出 + if len(inputs) > 1: + inputs = self._cell_backward_hook(inputs) + # 如果输入的数量小于1，则将输入转换为Cell的输出，并将输入转换为tuple + else: + inputs = self._cell_backward_hook(*inputs) + # 如果输入的类型是tuple，则调用construct方法，将输入转换为tuple + if isinstance(inputs, tuple): + outputs = self.construct(*inputs) + # 如果输入的类型不是tuple，则调用construct方法，将输入转换为tuple + else: + outputs = self.construct(inputs) + # 将输出转换为Cell的输出 + outputs = self._cell_backward_hook(outputs) + # 返回输出 + return outputs + + def register_backward_hook(self, hook_fn): + """ + Register the backward hook function. + + Note: + - The `register_backward_hook(hook_fn)` does not work in graph mode or ms_function. + - The 'hook_fn' must be defined as the following code. + `cell_id` is the information of registered Cell object, including name and ID. `grad_input` is the + gradient passed to the Cell. `grad_output` is the gradient computed and passed to the next Cell or + primitive, which may be modified by returning a new output gradient. + - The 'hook_fn' should have the following signature: + hook_fn(cell_id, grad_input, grad_output) -> New output gradient or none. + - The 'hook_fn' is executed in the python environment. In order to prevent running failed when switching to + graph mode, it is not recommended to write it in the `construct` function of Cell object. In the pynative + mode, if the `register_backward_hook` function is called in the `construct` function of the Cell object, + a hook function will be added at each run time of Cell object. + + Args: + hook_fn (function): Python function. Backward hook function. + + Returns: + Handle, it is an instance of `mindspore.common.hook_handle.HookHandle` and corresponding to the `hook_fn` . + The handle can be used to remove the added `hook_fn` by calling `handle.remove()` . + + Raises: + TypeError: If the `hook_fn` is not a function of python. + + Supported Platforms: + ``Ascend`` ``GPU`` ``CPU`` + + Examples: + >>> import numpy as np + >>> import mindspore + >>> import mindspore.nn as nn + >>> from mindspore import Tensor + >>> from mindspore import context + >>> from mindspore.ops import GradOperation + >>> context.set_context(mode=context.PYNATIVE_MODE) + >>> def backward_hook_fn(cell_id, grad_input, grad_output): + ... print("backward input: ", grad_input) + ... print("backward output: ", grad_output) + ... + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.relu = nn.ReLU() + ... self.handle = self.relu.register_backward_hook(backward_hook_fn) + ... + ... def construct(self, x): + ... x = x + x + ... x = self.relu(x) + ... return x + >>> grad = GradOperation(get_all=True) + >>> net = Net() + >>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32))) + backward input: (Tensor(shape=[1], dtype=Float32, value= [ 1.00000000e+00]),) + backward output: (Tensor(shape=[1], dtype=Float32, value= [ 1.00000000e+00]),) + >>> print(output) + (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),) + """ + if context.get_context("mode") != context.PYNATIVE_MODE: + logger.warning(f"'register_backward_hook' function is only supported in pynative mode, you can use " + f"context.set_context to set pynative mode.") + return HookHandle() + + if not isinstance(hook_fn, (FunctionType, MethodType)): + raise TypeError(f"When using'register_backward_hook(hook_fn)', the type of 'hook_fn' should be python " + f"function, but got {type(hook_fn)}.") + # 判断hook_fn是否为函数 + if self._cell_backward_hook is None: + # 如果没有设置backward hook，则设置backward hook + self._enable_backward_hook = True + # 设置backward hook为True + self._cell_backward_hook = inner.CellBackwardHook(self.cls_name + "(" + str(id(self)) + ")") + # 创建cell_backward_hook + backward_hook_key = self._cell_backward_hook.register_backward_hook(hook_fn) + # 注册backward hook + handle = HookHandle(self, backward_hook_key, "_cell_backward_hook") + else: + # 如果设置了backward hook，则注册backward hook + backward_hook_key = self._cell_backward_hook.register_backward_hook(hook_fn) + # 注册backward hook + handle = HookHandle(self, backward_hook_key, "_cell_backward_hook") + # 返回handle + return handle - def set_param_ps(self, recurse=True, init_in_server=False): + def set_param_ps(self, recurse=True, init_in_server=False):#设置可训练的参数是否由参数服务器进行更新，以及是否由可训练的参数在服务器上进行初始化 """ Set whether the trainable parameters are updated by parameter server and whether the trainable parameters are initialized on server. @@ -1158,10 +2375,23 @@ class Cell(Cell_): for param in params: param.set_param_ps(init_in_server) - def set_comm_fusion(self, fusion_type, recurse=True): + def set_param_fl(self, push_to_server=False, pull_from_server=False, requires_aggr=True):#设置参数和服务器交互的方式 """ - Set `comm_fusion` for all the parameters in the Net. Please refer to the description of - `mindspore.common.parameter.comm_fusion`. + Set the way of parameter and server interaction. + + Args: + push_to_server (bool): Whether the parameter should be pushed to server. Default: False. + pull_from_server (bool): Whether the parameter should be pulled from server. Default: False. + requires_aggr (bool): Whether the parameter should be aggregated in the server. Default: True. + """ + params = self.parameters_and_names() + for param in params: + param[1].set_param_fl(push_to_server, pull_from_server, requires_aggr) + + def set_comm_fusion(self, fusion_type, recurse=True):#为cell中每个参数设置 `comm_fusion` + """ + Set `comm_fusion` for all the parameters in this cell. Please refer to the description of + :class:`mindspore.Parameter.comm_fusion`. Note: The value of attribute will be overwritten when the function is called multiply. @@ -1177,76 +2407,236 @@ class Cell(Cell_): def _set_recompute_scope(self, mode): prefix = 'recompute_' - if mode is True: + # 如果模式存在，则设置_scope为prefix + if mode: if self._scope is None: self._scope = prefix + # 如果_scope不以prefix开头，则更新_scope elif not self._scope.startswith(prefix): self._scope = prefix + self._scope - elif not self._scope is None and self._scope.startswith(prefix): + # 如果_scope以prefix开头，则从_scope中删除prefix + elif self._scope is not None and self._scope.startswith(prefix): self._scope = self._scope[len(prefix):] - def recompute(self, mode=True): + def _mp_comm_recompute(self, mp_comm_recompute=True): """ - Set the cell recomputed. All the primitive in the cell will be set recomputed. If a primitive - set recomputed feeds into some backward nodes for computing gradient, rather than storing the - intermediate activation computed in forward pass, we will recompute it in backward pass. + Set the model parallel communication in cell recomputed. + """ + # 遍历每一个primitive + for _, value in self._primitives.items(): + # 如果primitive存在 + if value: + # 将recompute_comm_op设置为mp_comm_recompute + value.add_prim_attr("recompute_comm_op", mp_comm_recompute) + # 遍历每一个cell + for cell in self.cells(): + # 调用cell的_mp_comm_recompute函数 + cell._mp_comm_recompute(mp_comm_recompute) - Note: + def _parallel_optimizer_comm_recompute(self, parallel_optimizer_comm_recompute=False): + """ + Set the parallel optimizer communication in cell recomputed. + """ + for param in self.trainable_params(): + param.parallel_optimizer_comm_recompute = parallel_optimizer_comm_recompute - - If the computation involves something like randomization or global variable, the equivalence - is not guaranteed currently. - - If the recompute api of a primitive in this cell is also called, the recompute mode of this - primitive is subject to the recompute api of the primitive. + def _recompute_slice_activation(self, slice_activation=False): + """ + Slice the cell output which would remains in memory. + """ + # 遍历每一个primitive + for _, value in self._primitives.items(): + # 如果primitive存在 + if value: + # 将slice_activation设置为True + value.add_prim_attr("slice_activation", slice_activation) + # 遍历每一个cell + for cell in self.cells(): + # 调用cell的_recompute_slice_activation函数 + cell._recompute_slice_activation(slice_activation) - Args: - mode (bool): Specifies whether the cell is recomputed. Default: True. + def _recompute(self, mode=True, output_recompute=False): + """ + Set the cell recomputed. """ - if context.get_context("mode") == context.PYNATIVE_MODE: - raise TypeError("Recompute is not supported in pynative mode currently.") + if context._get_mode() == context.PYNATIVE_MODE: + # 如果当前模式为PYNATIVE_MODE，则抛出异常 + raise TypeError("Recompute is not supported in pynative mode currently, you can use " + "'context.set_context(mode=context.GRAPH_MODE)' to set graph mode.") + # 检查mode是否为布尔值 Validator.check_bool(mode) + # 检查output_recompute是否为布尔值 + Validator.check_bool(output_recompute) + # 如果_has_config_recompute为False，则设置_has_config_recompute为True + if not self._has_config_recompute: + self._has_config_recompute = True + # 否则抛出异常 + else: + raise RuntimeError("The recompute interface can be configured only once." + " When the parent cell is configured, the child cell should not be configured") + # 设置_recompute_scope self._set_recompute_scope(mode) + # 如果mode为True，且output_recompute为False，则添加output_no_recompute标志 + if mode and not output_recompute: + self.add_flags(output_no_recompute=True) + # 遍历所有的cell for cell in self.cells(): - cell.recompute(mode) - - -class GraphKernel(Cell): - """ - Base class for GraphKernel. - - A `GraphKernel` a composite of basic primitives and can be compiled into a fused kernel automatically when - enable_graph_kernel in context is set to True. - - This class is deprecated from version 1.3 and will be removed in a future version, use Cell instead. - - GraphKernel is not supported user-defined cells anymore, the `GraphKernel` objects will be treated as - normal `Cell` objects. + # 调用cell的_recompute函数 + cell._recompute(mode, True) + @args_type_check(mp_comm_recompute=bool, parallel_optimizer_comm_recompute=bool) + def recompute(self, **kwargs):#cell中recompute的设置 + """ + Set the cell recomputed. All the primitive in the cell except the outputs will be set recomputed. + If a primitive set recomputed feeds into some backward nodes for computing gradient, rather than + storing the intermediate activation computed in forward pass, we will recompute it in backward pass. - Args: - auto_prefix (bool): Recursively generate namespaces. Default: True. - flags (dict) : Set graph flags. Default: None. + Note: - Supported Platforms: - ``Ascend`` ``GPU`` + - If the computation involves something like randomization or global variable, the equivalence + is not guaranteed currently. + - If the recompute api of a primitive in this cell is also called, the recompute mode of this + primitive is subject to the recompute api of the primitive. + - The interface can be configured only once. + Therefore, when the parent cell is configured, the child cell should not be configured. + - The outputs of cell are excluded from recomputation by default, which is based on our configuration + experience to reduce memory footprint. If a cell has only one primitive and the primitive is wanted + to be set recomputed, use the recompute api of the primtive. + - When the memory remains after applying the recomputation, configuring 'mp_comm_recompute=False' + to improve performance if necessary. + - When the memory still not enough after applying the recompute, configuring + 'parallel_optimizer_comm_recompute=True' to save more memory if necessary. + Cells in the same fusion group should have the same parallel_optimizer_comm_recompute configures. - Examples: - >>> class Relu(nn.GraphKernel): - ... def __init__(self): - ... super(Relu, self).__init__() - ... self.max = P.Maximum() - ... - ... def construct(self, x): - ... return self.max(P.Fill()(P.DType()(x), P.Shape()(x), 0.0), x) - """ + Args: + mp_comm_recompute (bool): Specifies whether the model parallel communication operators + in the cell are recomputed in auto parallel or semi auto parallel mode. Default: True. + parallel_optimizer_comm_recompute (bool): Specifies whether the communication operator allgathers + introduced by optimizer shard are recomputed in auto parallel or semi auto parallel mode. + Default: False. + """ + self._recompute() + # 如果kwargs中有mp_comm_recompute，则调用_mp_comm_recompute函数 + if'mp_comm_recompute' in kwargs.keys(): + self._mp_comm_recompute(kwargs.get('mp_comm_recompute', False)) + # 如果kwargs中有parallel_optimizer_comm_recompute，则调用_parallel_optimizer_comm_recompute函数 + if 'parallel_optimizer_comm_recompute' in kwargs.keys(): + if (kwargs.get('parallel_optimizer_comm_recompute', False) and + context.get_auto_parallel_context("pipeline_stages") > 1): + logger.warning("Currently, the communication operator allgathers introduced by optimizer shard " + "are not support recomputation in pipeline parallel.") + elif context.get_auto_parallel_context("pipeline_stages") == 1: + self._parallel_optimizer_comm_recompute(kwargs.get('parallel_optimizer_comm_recompute', False)) + # 如果kwargs中有recompute_slice_activation，则调用_recompute_slice_activation函数 + if'recompute_slice_activation' in kwargs.keys(): + self._recompute_slice_activation(kwargs.get('recompute_slice_activation', False)) + + # 遍历kwargs中的key，如果key不在'recompute'中，则抛出ValueError异常 + for key, _ in kwargs.items(): + if key not in ('mp_comm_recompute', 'parallel_optimizer_comm_recompute', 'recompute_slice_activation'): + raise ValueError("For 'recompute', keyword '%s' is not recognized! " + "the key kwargs must be 'mp_comm_recompute', " + "'parallel_optimizer_comm_recompute', 'recompute_slice_activation'" % key) + + def infer_param_pipeline_stage(self):#静态分析pipeline_stages中的数据 + """ + Infer pipeline stages of all parameters in the cell. - @deprecated("1.3", "Cell", True) - def __init__(self, auto_prefix=True, flags=None): - super(GraphKernel, self).__init__(auto_prefix, flags) + Note: + - If a parameter does not belong to any cell which has been set pipeline_stage, + the parameter should use add_pipeline_stage to add it's pipeline_stage information. + - If a parameter P has been used by two operators in different stages "stageA" and "stageB", + the parameter P should use P.add_pipeline_stage(stageA) and P.add_pipeline_stage(stageB) + to add it's stage information before using infer_param_pipeline_stage. - def construct(self): - raise NotImplementedError + Returns: + The params belong to current stage in pipeline parallel. + Raises: + RuntimeError: If there is a parameter does not belong to any stage. + """ + from mindspore.parallel._utils import _get_global_rank, _get_device_num + logger.warning(f"This interface may be deleted in the future.") + # 获取自动并行上下文中的pipeline_stages + stage_num = context.get_auto_parallel_context("pipeline_stages") + # 获取设备数量 + device_num = _get_device_num() + # 获取当前程序的计算节点 + rank_id = _get_global_rank() + # 计算每个stage的设备数量 + per_stage_devices = device_num // stage_num + # 获取当前stage的索引 + current_stage = rank_id // per_stage_devices + # 初始化参数列表 + params = [] + # 遍历训练参数 + for param in self.trainable_params(): + # 如果参数不在任何stage中，抛出异常 + if not param._pipeline_stage_list: + raise RuntimeError("For 'infer_param_pipeline_stage', the parameter {} does not belong to any stage, " + "please check whether the cell where the param locates has been set " + "'pipeline_stage'. Otherwise, the parameter should use 'add_pipeline_stage' " + "to add its stage information".format(param.name)) + # 如果当前stage在参数的pipeline_stage列表中，则将参数添加到参数列表中 + if current_stage in param._pipeline_stage_list: + params.append(param) + return params + + #检查输入的是否与动态形状输入的一致 + def _check_compile_dynamic_shape(self, *inputs): + """ + Check if graph has been compiled with dynamic shape. -class GraphCell(Cell): + Args: + inputs (tuple): Inputs of the Cell object. + """ + # 获取输入的长度 + len_inputs = len(inputs) + # 获取动态形状输入的长度 + len_dynamic_shape_inputs = len(self._dynamic_shape_inputs) + # 检查输入的长度是否与动态形状输入的长度一致 + if len_dynamic_shape_inputs!= len_inputs: + raise ValueError( + f"For 'set_inputs', the Length of Tensor should be {len_inputs}, but got {len_dynamic_shape_inputs}." + f"For'set_inputs', the Length of Tensor should be {len_inputs}, but got {len_dynamic_shape_inputs}." + ) + # 遍历动态形状输入 + for tensor_index in range(len_dynamic_shape_inputs): + # 获取动态形状输入 + i_dynamic_shape_inputs = self._dynamic_shape_inputs[tensor_index] + # 获取输入 + i_inputs = inputs[tensor_index] + # 检查输入的数据类型是否与动态形状输入的数据类型一致 + if i_dynamic_shape_inputs.dtype is not i_inputs.dtype: + raise TypeError( + f"For 'set_inputs', the DataType of Tensor should be {i_inputs.dtype}, but got " + f"For'set_inputs', the DataType of Tensor should be {i_inputs.dtype}, but got " + f"{i_dynamic_shape_inputs.dtype}." + ) + # 获取输入的形状 + set_inputs_shape = list(i_dynamic_shape_inputs.shape) + inputs_shape = list(i_inputs.shape) + # 检查输入的形状是否与动态形状输入的形状一致 + if len(inputs_shape)!= len(set_inputs_shape): + raise ValueError( + f"For 'set_inputs' the Dimension of Tensor shape must be {len(inputs_shape)}, but got " + f"For'set_inputs' the Dimension of Tensor shape must be {len(inputs_shape)}, but got " + f"{len(set_inputs_shape)}." + ) + # 遍历输入的形状 + for shape_index in i_dynamic_shape_inputs.shape: + # 检查形状是否为-1 + if shape_index!= -1: + # 获取动态形状输入的形状索引 + dynamic_index = i_dynamic_shape_inputs.shape.index(shape_index) + # 检查输入的形状是否与动态形状输入的形状一致 + if set_inputs_shape[dynamic_index]!= inputs_shape[dynamic_index]: + raise ValueError( + f"For 'Length of Tensor shape', the value must be the same with that of inputs, but" + f" got {i_dynamic_shape_inputs.shape}." + ) + + +class GraphCell(Cell):# GraphCell类 """ Base class for running the graph loaded from a MindIR. @@ -1254,7 +2644,16 @@ class GraphCell(Cell): diagram, and can only use data that shape and type are the same as the input when exporting the MindIR. Args: - graph (object): A compiled graph loaded from MindIR. + graph (FuncGraph): A compiled graph loaded from MindIR. + params_init (dict): Parameters need to be inited in the graph. + The key is the parameter name whose type is str, and the value is a Tensor or Parameter. + If the parameter exists in the graph according to the name, update it's value. + If the parameter does not exist, ignore it. Default: None. + Raises: + TypeError: If the `graph` is not a FuncGraph. + TypeError: If the `params_init` is not a dict. + TypeError: If the key of the `params_init` is not a str. + TypeError: If the value of the `params_init` is neither a Tensor nor a Parameter. Supported Platforms: ``Ascend`` ``GPU`` ``CPU`` @@ -1262,24 +2661,68 @@ class GraphCell(Cell): Examples: >>> import numpy as np >>> import mindspore.nn as nn - >>> from mindspore import Tensor - >>> from mindspore.train import export, load + >>> from mindspore import Tensor, export, load >>> - >>> net = nn.Conv2d(1, 1, kernel_size=3) + >>> net = nn.Conv2d(1, 1, kernel_size=3, weight_init="ones") >>> input = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32)) >>> export(net, input, file_name="net", file_format="MINDIR") >>> graph = load("net.mindir") >>> net = nn.GraphCell(graph) >>> output = net(input) + >>> print(output) + [[[[4. 6. 4.] + [6. 9. 6.] + [4. 6. 4.]]]] """ - def __init__(self, graph): + def __init__(self, graph, params_init=None): + # 初始化GraphCell类 super(GraphCell, self).__init__(auto_prefix=True) + # 判断传入的graph是否为FuncGraph类型 if not isinstance(graph, FuncGraph): - raise TypeError(f"graph must be a FuncGraph loaded from MindIR, but got {type(graph)}.") + raise TypeError(f"For 'GraphCell', the argument 'graph' must be a FuncGraph loaded from MindIR, " + f"but got type {type(graph)}.") self.graph = graph + # 初始化params_init参数 + params_init = {} if params_init is None else params_init + # 判断params_init参数是否为字典类型 + if not isinstance(params_init, dict): + raise TypeError(f"For 'GraphCell', the argument 'params_init' must be a dict, but got {type(params_init)}.") + # 遍历params_init参数，将其转换为Tensor类型 + for name, value in params_init.items(): + if not isinstance(name, str) or not isinstance(value, Tensor): + raise TypeError("For 'GraphCell', the key of the 'params_init' must be str, " + "and the value must be Tensor or Parameter, " + f"but got the key type: {type(name)}, and the value type: {type(value)}") + + # 更新graph的hyper_params + params_dict = update_func_graph_hyper_params(self.graph, params_init) + # 遍历params_dict，将其赋值给self._params + for name, param in params_dict.items(): + self._params[name] = param + def construct(self, *inputs): + # 返回graph函数调用 return self.graph(*inputs) def __call__(self, *inputs): + # 设置graph_load_from_mindir参数 + self.phase = "graph_load_from_mindir" + # 添加graph_load_from_mindir参数 + self._add_attr("graph_load_from_mindir", self.graph) + # 返回compile_and_run函数调用 return self.compile_and_run(*inputs) + + +def _check_param_list_tuple(value): + """ + Check the type of input in list or tuple is Parameter. + :param value: list or tuple. + :return: The types of all inputs are parameter. + """ + for item in value: + # 如果item不是Parameter类型，返回False + if not isinstance(item, Parameter): + return False + # 如果所有的item都是Parameter类型，返回True + return True