{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# 添加AtenToMindSpore算子映射关系教程\n",
    "\n",
    "`Linux` `Ascend` `GPU` `CPU` `模型迁移` `Pytorch` `ATen`\n",
    "\n",
    "[![](https://gitee.com/mindspore/docs/raw/master/resource/_static/logo_source.png)](https://gitee.com/mindspore/mindinsight/blob/master/ecosystem_tools/mindconverter/tutorial/add_aten2mindspore_operator_mapper_tutorial.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 概述\n",
    "\n",
    "MindConverter API 基于PyTorch源生计算图[TorchScript](https://pytorch.org/docs/stable/jit.html)进行模型迁移。通过将TorchScript中的[ATen算子](https://pytorch.org/cppdocs/index.html#aten)映射到MindSpore中的对应数学实现的算子，来生成MindSpore模型脚本和权重文件。因此ATen算子到MindSpore算子的映射关系的正确性是决定迁移结果正确与否的关键要素之一。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 环境准备\n",
    "\n",
    "本案例需要安装以下依赖库：\n",
    "\n",
    "```bash\n",
    "pip install mindspore==1.6.0\n",
    "pip install torch==1.8.1 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 自定义算子映射基类\n",
    "\n",
    "```python\n",
    "class AtenToMindSporeMapper(...):\n",
    "    \"\"\"实现ATen算子到MindSpore算子的基类。\"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def _operation_name_in_ms(*args, **kwargs) -> str:\n",
    "        \"\"\"返回MindSpore算子名。\"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def _convert_trainable_weights(**kwargs) -> dict:\n",
    "        \"\"\"返回由ATen算子中的权重信息转换后的MindSpore算子中的权重。\"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def _generate_snippet_template(**kwargs) -> (dict, dict, list, tuple):\n",
    "        \"\"\"\n",
    "        返回用于生成MindSpore模型脚本的信息。\n",
    "\n",
    "        Returns:\n",
    "            Dict, template.  # 分别保存MindSpore模型脚本中的`__init__`和`construct`部分的算子代码。\n",
    "                如果template =\n",
    "                {\n",
    "                    variable_slot: {\n",
    "                        init: [\"init_code_0\", \"init_code_1\",...],\n",
    "                        construct: [\"construct_code_0\", \"construct_code_1\",...]\n",
    "                    }\n",
    "                }，\n",
    "                那么，该算子在MindSpore模型脚本中的代码为：\n",
    "                class A(nn.Cell):\n",
    "                    def __init__(self):\n",
    "                        super(A, self).__init__()\n",
    "                        init_code_0\n",
    "                        init_code_1\n",
    "                        ...\n",
    "\n",
    "                    def construct(self, x):\n",
    "                        construct_code_0\n",
    "                        construct_code_1\n",
    "                        ...\n",
    "            Dict, exchange_msg.  # 保存来自于其他算子的与该算子有关的信息，例如：该算子的输入来自于其他算子的输出。\n",
    "                exchange_msg =\n",
    "                {\n",
    "                    variable_slot: {\n",
    "                        operation: op_name,  # MindSpore算子名。\n",
    "                        variable_name: None,  # 算子在MindSpore模型脚本中的引用变量名，由后续模块生成，此处设置为None。\n",
    "                        output_type: \"tensor\",  # 该算子的输出类型，通常设置为张量即可。\n",
    "                        inputs: [],  # 该算子的真实输入，由后续模块生成，此处设置为list()。\n",
    "                        args: args,  # 该算子的参数。\n",
    "                        weights: weights,  # 该算子包含的全部张量，包括静态张量与权重张量。\n",
    "                        trainable_params: trainable_params,  # 该算子的权重信息，后续会保存在CheckPoint file中。\n",
    "                        parameters_declared: parameters_declared,  # 与trainable_params配合使用，在MindSpore模型脚本中加载非算子权重的静态张量。该字段可能不存在。\n",
    "                        group_inputs: group_inputs  # 与inputs配合使用，划分该算子来自于其他算子输出的输入变量。该字段可能不存在。\n",
    "                    }\n",
    "                }\n",
    "            List, outputs_list.  # 该算子的输出名。算子可能存在多个输出，因此用列表封装。\n",
    "            Tuple, outputs_mapping.  # ATen算子与MindSpore算子的输出的对应关系，主要用于都输出算子的映射。\n",
    "        \"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def get_args_name_list(**kwargs) -> list:\n",
    "        \"\"\"\n",
    "        返回有效的ATen算子参数列表。\n",
    "        同一个ATen算子可能存在多种声明，通过参数个数和参数类型进行区分，从而适配前端算子的灵活表达。\n",
    "        以`aten::zeros`为例：\n",
    "            at::Tensor at::zeros(at::IntArrayRef size, at::TensorOptions options = {})\n",
    "            at::Tensor at::zeros(at::IntArrayRef size, c10::optional<at::DimnameList> names, at::TensorOptions options = {})\n",
    "            at::Tensor at::zeros(at::IntArrayRef size, c10::optional<at::ScalarType> dtype, c10::optional<at::Layout> layout, c10::optional<at::Device> device, c10::optional<bool> pin_memory)\n",
    "            at::Tensor at::zeros(at::IntArrayRef size, c10::optional<at::DimnameList> names, c10::optional<at::ScalarType> dtype, c10::optional<at::Layout> layout, c10::optional<at::Device> device, c10::optional<bool> pin_memory)\n",
    "        那么，kwargs[\"args_name\"] =\n",
    "            {\n",
    "                2: ['shape', 'unused'],\n",
    "                3: ['shape', 'unused', 'unused'],\n",
    "                5: ['shape', 'dtype', 'unused', 'unused', 'unused'],\n",
    "                6: ['shape', 'unused', 'dtype', 'unused', 'unused', 'unused']\n",
    "            }\n",
    "\n",
    "        Returns:\n",
    "            List, args_name_list.  # 通过参数数量确定的ATen算子参数名列表。\n",
    "        \"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def _params_parser(raw_params, args_name, trainable_params) -> (list, dict, list):\n",
    "        \"\"\"\n",
    "        对ATen算子的参数进行解析，获取用于生成MindSpore对应算子的信息。\n",
    "\n",
    "        Args:\n",
    "            raw_params (dict): ATen算子的参数。\n",
    "            args_name (Union<list, dict>): ATen算子的预设参数名列表。仅有一个声明时为列表(list)，有多个声明时为字典(dict)。\n",
    "            trainable_params (dict): ATen算子的张量信息。\n",
    "\n",
    "        Returns:\n",
    "            list, inputs.  # ATen算子的参数在MindSpore模型脚本中的对应表达形式。\n",
    "            dict, args.  # ATen算子中参数的值。\n",
    "            list, group_inputs.  # 与`inputs`对应，用于正确地生成MindSpore算子中来自于其他算子输出的输入变量。\n",
    "        \"\"\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 自定义添加算子映射关系\n",
    "\n",
    "确定ATen算子和MindSpore算子的对应关系，通常需要对于算子的数学实现有一定的了解。\n",
    "\n",
    "|算子类型|信息来源|说明|\n",
    "|:-----:|:-----:|:-----:|\n",
    "|`aten`|[API文档](https://pytorch.org/cppdocs/index.html#aten)<br>[ATen算子注册表](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml)|优先使用API文档中的`Functions`和`Class Tensor`中的声明。|\n",
    "|`mindspore`|[API文档](https://www.mindspore.cn/docs/zh-CN/master/index.html)|MindSpore网络中所需算子主要在`mindspore.nn`, `mindspore.ops`和`mindspore.numpy`中。|\n",
    "\n",
    "通过查询ATen算子的API文档，获取`aten::zeros`的声明：\n",
    "\n",
    "```text\n",
    "at::Tensor at::zeros(at::IntArrayRef size, at::TensorOptions options = {})\n",
    "at::Tensor at::zeros(at::IntArrayRef size, c10::optional<at::DimnameList> names, at::TensorOptions options = {})\n",
    "at::Tensor at::zeros(at::IntArrayRef size, c10::optional<at::ScalarType> dtype, c10::optional<at::Layout> layout, c10::optional<at::Device> device, c10::optional<bool> pin_memory)\n",
    "at::Tensor at::zeros(at::IntArrayRef size, c10::optional<at::DimnameList> names, c10::optional<at::ScalarType> dtype, c10::optional<at::Layout> layout, c10::optional<at::Device> device, c10::optional<bool> pin_memory)\n",
    "```\n",
    "\n",
    "通过查询MindSpore算子的API文档，获取具有相似数学实现的MindSpore算子`mindspore.numpy.zeros`的声明：\n",
    "\n",
    "```text\n",
    "mindspore.numpy.zeros(shape, dtype=mstype.float32)\n",
    "    Returns a new tensor of given shape and type, filled with zeros.\n",
    "\n",
    "    Parameters\n",
    "        shape(Union[int, tuple, list]) - the shape of the new tensor.\n",
    "        dtype(Union[mindspore.dtype, str], optional) - Designated tensor dtype. Default is mstype.float32.\n",
    "\n",
    "    Returns\n",
    "        Tensor, with the designated shape and dtype, filled with zeros.\n",
    "\n",
    "    Raises\n",
    "        TypeError - If input arguments have types not specified above.\n",
    "        ValueError - If shape entries have values < 0.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 根据`aten::zeros`和`mindspore.numpy.zeros`的声明实现映射。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# 固定导入，多数情况下不需要修改。\n",
    "from mindconverter.graph_based_converter.common.utils import reset_template_and_exchange_msg\n",
    "from mindconverter.graph_based_converter.constant import WeightType, PYTORCH_MS_MAP\n",
    "from mindconverter.graph_based_converter.mapper.base import AtenToMindSporeMapper\n",
    "\n",
    "\n",
    "class ZerosMapper(AtenToMindSporeMapper):\n",
    "    \"\"\"定义映射关系类，该类名与ATen算子保持一致。\"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def _operation_name_in_ms(*args, **kwargs):\n",
    "        \"\"\"返回与ATen算子对应的MindSpore算子名。其中的`mindspore.numpy`已重命名为`ms_np`，不在此处返回。\"\"\"\n",
    "        return \"zeros\"\n",
    "\n",
    "    @staticmethod\n",
    "    def _convert_trained_weights(**kwargs):\n",
    "        \"\"\"返回该算子的张量信息。\"\"\"\n",
    "        weights = kwargs.get(\"weights\", list())\n",
    "        args_name = {\n",
    "            2: [\"shape\", \"unused\"],\n",
    "            3: [\"shape\", \"unused\", \"unused\"],\n",
    "            5: [\"shape\", \"dtype\", \"unused\", \"unused\", \"unused\"],\n",
    "            6: [\"shape\", \"unused\", \"dtype\", \"unused\", \"unused\", \"unused\"]\n",
    "        }\n",
    "        # 获取该算子的参数名列表。\n",
    "        args_name_list = ZerosMapper.get_args_name_list(**kwargs, args_name=args_name)\n",
    "        # 获取该算子张量信息。\n",
    "        trainable_params = dict()\n",
    "        for weight in weights:\n",
    "            trainable_params[args_name_list[weight.location]] = {\n",
    "                \"data\": weight.value, \"location\": weight.location,\n",
    "                \"type\": WeightType.PARAMETER.value, \"onnx_name\": weight.name\n",
    "            }\n",
    "        return trainable_params\n",
    "\n",
    "    @staticmethod\n",
    "    def _generate_snippet_template(**kwargs):\n",
    "        \"\"\"返回用于生成MindSpore模型脚本中的该算子代码所需要的信息。\"\"\"\n",
    "        # 固定代码段，不需要修改。\n",
    "        template, exchange_msg, outputs_list, outputs_mapping = AtenToMindSporeMapper._generate_snippet_template(\n",
    "            **kwargs)\n",
    "        raw_params = kwargs.get(\"raw_params\")\n",
    "        if not raw_params:\n",
    "            return template, exchange_msg, outputs_mapping, outputs_mapping\n",
    "\n",
    "        op = kwargs.get(\"operation\")\n",
    "\n",
    "        variable_slot = \"var_0\"\n",
    "        trainable_params = kwargs.get(\"trainable_params\", dict())\n",
    "\n",
    "        # ATen算子的参数名配置。\n",
    "        args_name = {\n",
    "            2: [\"shape\", \"unused\"],\n",
    "            3: [\"shape\", \"unused\", \"unused\"],\n",
    "            5: [\"shape\", \"dtype\", \"unused\", \"unused\", \"unused\"],\n",
    "            6: [\"shape\", \"unused\", \"dtype\", \"unused\", \"unused\", \"unused\"]\n",
    "        }\n",
    "\n",
    "        # 获取算子对应的参数列表。\n",
    "        # ATen算子中的`shape`参数是一个元组，该元组可能是静态配置，也可能是来自于其他算子输出的动态输入。\n",
    "        # 当其为动态输入时，则`shape = (shape0, shape1, shape2,...)`。\n",
    "        # if return_raw:\n",
    "        #     return [(shape0, shape1, shape2,...), \"unused\"]\n",
    "        # else:\n",
    "        #     return [shape0, shape1, shape2,..., \"unused\"]\n",
    "        raw_args_name_list = ZerosMapper.get_args_name_list(params=raw_params, args_name=args_name, return_raw=True)\n",
    "\n",
    "        # 解析ATen算子参数，获取MindSpore算子对应的表达。\n",
    "        inputs, args, group_inputs = ZerosMapper._params_parser(raw_params, args_name, trainable_params)\n",
    "\n",
    "        # 对于某些从ATen算子中获取的参数，无法直接用于MindSpore算子中，需要进行相应的处理。\n",
    "        dtype = args.get(\"dtype\")\n",
    "        args[\"dtype\"] = PYTORCH_MS_MAP[\"default\"] if dtype is None else PYTORCH_MS_MAP[dtype]\n",
    "\n",
    "        # 该算子在`__init__`中的代码。\n",
    "        init_template_list = [f\"self.{{{variable_slot}}}_{arg_name} = {{{arg_name}}}\" for arg_name in args]\n",
    "\n",
    "        # 获取`parameters_declared`，多数情况下为固定代码，不需要修改。\n",
    "        parameters_declared = ZerosMapper.generate_parameters_declared(variable_slot, init_template_list, args,\n",
    "                                                                       trainable_params)\n",
    "        # 根据`raw_args_name_list和`inputs`的长度关系来确定生成`shape`还是`(shape0, shape1, shape2,...)`。\n",
    "        diff = len(inputs) - len(raw_args_name_list)\n",
    "        if diff > 0:\n",
    "            construct_template = f\"opt_{{{variable_slot}}} = ms_np.{op}\" \\\n",
    "                                 f\"(({', '.join(inputs[:diff + 1])}), self.{{{variable_slot}}}_dtype)\"\n",
    "        else:\n",
    "            construct_template = f\"opt_{{{variable_slot}}} = ms_np.{op}\" \\\n",
    "                                 f\"({inputs[0]}, self.{{{variable_slot}}}_dtype)\"\n",
    "        # 生成`template`和`exchange_msg`。多数情况下为固定代码，不需要修改。\n",
    "        template, exchange_msg = reset_template_and_exchange_msg(template, exchange_msg, variable_slot,\n",
    "                                                                 init_template_list, [construct_template], args,\n",
    "                                                                 trainable_params, parameters_declared, group_inputs)\n",
    "        return template, exchange_msg, outputs_list, outputs_mapping"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 新增映射关系文件到工具中，使该映射关系在工具运行中生效。\n",
    "\n",
    "- 将该映射关系脚本命名为`zeros_mapper`，该文件命名以ATen算子为准。\n",
    "- 将该映射关系文件放入`mindconverter/graph_based_converter/mapper/aten/prim`目录下，`aten`当中有`nn`，`ops`，`prim`目录，分别对应MindSpore算子的`nn`，`ops`，`numpy`层算子。\n",
    "- 在[映射关系注册表](../mindconverter/graph_based_converter/mapper/aten_to_ms.json)中增加`\"aten::zeros\": \"mindconverter.graph_based_converter.mapper.aten.prim.zeros_mapper.ZerosMapper\"`，\n",
    "\n",
    "用于动态引用该映射关系文件进行算子转换。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 验证算子映射关系的正确性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from mindconverter.graph_based_converter.mapper.base import AtenToMindSporeMapper\n",
    "from mindconverter.graph_based_converter.common.code_fragment import Fragment\n",
    "\n",
    "\n",
    "def test_mapper(aten_info, inputs_code):\n",
    "    \"\"\"\n",
    "    Test mapper.\n",
    "\n",
    "    Args:\n",
    "        aten_info (dict): ATen operator info.\n",
    "        {\n",
    "            'op_name': op_name,\n",
    "            'attribute': dict(),\n",
    "            'weights': [NodeWeight(),...]\n",
    "        }\n",
    "    \"\"\"\n",
    "    template, exchange_msg, outputs_list, outputs_mapping = AtenToMindSporeMapper.convert(aten_info[\"op_name\"],\n",
    "                                                                                          aten_info[\"attribute\"],\n",
    "                                                                                          aten_info[\"weights\"])\n",
    "    exchange_msg['var_0']['variable_name'] = 'ms_zeros'\n",
    "    exchange_msg['var_0']['inputs'] = inputs_code\n",
    "\n",
    "    fragment = Fragment(data_entity=exchange_msg, code_template=template, outputs=outputs_list,\n",
    "                        outputs_mapping=outputs_mapping)\n",
    "\n",
    "    code = fragment()\n",
    "    init_code = code[0]\n",
    "    construct_code = code[1]\n",
    "    print('-' * 30, 'init_code', '-' * 30)\n",
    "    print('\\n'.join(init_code))\n",
    "    print('-' * 30, 'construct_code', '-' * 30)\n",
    "    print('\\n'.join(construct_code))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------------------------------ init_code ------------------------------\n",
      "self.ms_zeros_dtype = mindspore.int64\n",
      "------------------------------ construct_code ------------------------------\n",
      "opt_ms_zeros = ms_np.zeros(x, self.ms_zeros_dtype)\n"
     ]
    }
   ],
   "source": [
    "aten_info = {\n",
    "    \"op_name\": \"aten::zeros\",\n",
    "    \"attribute\": {\n",
    "        \"constant_0\": \"from_input\",\n",
    "        \"constant_1\": 4,\n",
    "        \"constant_2\": None,\n",
    "        \"constant_3\": None,\n",
    "        \"constant_4\": None,\n",
    "        \"input_shape\": -1,\n",
    "        \"output_shape\": (1, 3, 224, 224)},\n",
    "    \"weights\": list()\n",
    "}\n",
    "\n",
    "test_mapper(aten_info, inputs_code=[\"x\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------------------------------ init_code ------------------------------\n",
      "self.ms_zeros_shape1 = 3\n",
      "self.ms_zeros_shape3 = 224\n",
      "self.ms_zeros_dtype = mindspore.int64\n",
      "------------------------------ construct_code ------------------------------\n",
      "opt_ms_zeros = ms_np.zeros((x, self.ms_zeros_shape1, y, self.ms_zeros_shape3), self.ms_zeros_dtype)\n"
     ]
    }
   ],
   "source": [
    "aten_info = {\n",
    "    \"op_name\": \"aten::zeros\",\n",
    "    \"attribute\": {\n",
    "        \"constant_0[0(*,_,_,_)]\": \"from_input\",\n",
    "        \"constant_1[0(_,*,_,_)]\": 3,\n",
    "        \"constant_2[0(_,_,*,_)]\": \"from_input\",\n",
    "        \"constant_3[0(_,_,_,*)]\": 224,\n",
    "        \"constant_4\": None,\n",
    "        \"constant_5\": 4,\n",
    "        \"constant_6\": None,\n",
    "        \"constant_7\": None,\n",
    "        \"constant_8\": None,\n",
    "        \"input_shape\": -1,\n",
    "        \"output_shape\": (1, 3, 224, 224)\n",
    "    },\n",
    "    \"weights\": list()\n",
    "}\n",
    "\n",
    "test_mapper(aten_info, inputs_code=[\"x\", \"y\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}