pytorch_op_mapper.py

#   Copyright (c) 2019  PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import numpy as np
from x2paddle.core.op_mapper import OpMapper
from x2paddle.core.util import *
from x2paddle.core.program import PaddleGraph
from x2paddle.op_mapper.pytorch2paddle import prim
from x2paddle.op_mapper.pytorch2paddle import aten


class PyTorchOpMapper(OpMapper):
    def __init__(self, decoder):
        super(PyTorchOpMapper, self).__init__()
        self.script = decoder.script
        self.paddle_params = dict()
        self.outputs_info = {}  # key为output unique id，value为当前节点的输出名字
        self.pytorch_params = {}  # key为节点名，value为参数
        self.attrs = {}  # key为节点名，value为属性值
        self.output_index = 0
        self.dygraph_name_id = {}  # 动态图__init__输出名字中的id，key为kernel类型，value为id
        # 转换
        self.check_op(decoder.graph)
        self.graph, _ = self.traverse(decoder.graph)

    def check_op(self, script_graph):
        def _update_op_list(graph):
            for node in graph.nodes():
                op_list.append(node.kind())
                for block in node.blocks():
                    _update_op_list(block)

        op_list = list()
        _update_op_list(script_graph)
        op_list = list(set(op_list))
        unsupported_op_list = []
        for op in op_list:
            func_name = op.replace('::', '_')
            if not (hasattr(prim, func_name) or hasattr(aten, func_name)):
                unsupported_op_list.append(op)
        if len(unsupported_op_list) > 0:
            raise Exception("The kind {} in model is not supported yet.".format(
                unsupported_op_list))

    def traverse(self, script_graph, parent_layer=None):
        # 用于获取graph的输入
        def _update_graph_inputs(inputs, outputs):
            current_node_outputs.extend(outputs)
            for name in inputs:
                if name not in current_node_outputs:
                    graph_inputs.append(name)

        # 初始化
        graph = PaddleGraph(parent_layer)
        current_node_outputs = []
        graph_inputs = []
        # 转换输入节点
        if isinstance(script_graph, torch._C.Graph):
            for i, ivalue in enumerate(script_graph.inputs()):
                node = ivalue.node()
                if str(ivalue.type()) != "Tensor":
                    graph.set_name(str(ivalue.type()).split(".")[-1])
                inputs, outputs = self.data(graph, node, ivalue.unique())
        # 转换中间节点
        for node in script_graph.nodes():
            kind = node.kind()
            func_name = kind.replace('::', '_')
            if hasattr(prim, func_name):
                func = getattr(prim, func_name)
                inputs, outputs = func(self, graph, node)
                _update_graph_inputs(inputs, outputs)
            elif hasattr(aten, func_name):
                func = getattr(aten, func_name)
                inputs, outputs = func(self, graph, node)
                _update_graph_inputs(inputs, outputs)

        # 转换输出节点
        if hasattr(script_graph, 'returnNode'):
            for i, ivalue in enumerate(script_graph.returnNode().inputs()):
                if parent_layer.kernel == "prim.loop" and i == 0:
                    continue
                node = ivalue.node()
                script_unique_id = ivalue.unique()
                inputs, outputs = self.equal(
                    graph,
                    node,
                    uid=script_unique_id,
                    parent_layer=parent_layer,
                    index=i)
                _update_graph_inputs(inputs, outputs)
        # 设置graph的参数
        if isinstance(script_graph, torch._C.Graph):
            graph.set_parameters(self.paddle_params)
        return graph, graph_inputs

    def _get_outputs_name(self, node):
        outputs_name = []
        for output_ivalue in node.outputs():
            output_name = 'x' + str(self.output_index)
            script_unique_id = output_ivalue.unique()
            if script_unique_id in self.outputs_info:
                output_name = self.outputs_info[script_unique_id]
            self.outputs_info[script_unique_id] = output_name
            self.output_index += 1
            outputs_name.append(output_name)
        # if或loop节点没有输出的情况
        if len(list(node.outputs())) == 0:
            output_name = '_x' + str(self.output_index)
            self.output_index += 1
            outputs_name.append(output_name)
        return outputs_name

    def _check_input(self,
                     graph,
                     node,
                     output_name,
                     node_outputs,
                     add_dim=False):
        if node.kind() == "prim::GetAttr":
            param = self.pytorch_params[output_name]
            if isinstance(param, np.ndarray):
                if add_dim:
                    param = param[np.newaxis, :]
                self.paddle_params[output_name] = param
                graph.add_layer(
                    "fluid.dygraph.base.to_variable",
                    inputs={},
                    outputs=[output_name],
                    value="params[{}]".format(string(output_name)))
            else:
                if isinstance(param, dict) and "Tensor" in param:
                    graph.add_layer(
                        "prim.constant",
                        inputs={},
                        outputs=[output_name],
                        value=param["Tensor"])
                else:
                    graph.add_layer(
                        "prim.constant",
                        inputs={},
                        outputs=[output_name],
                        value=string(param)
                        if isinstance(param, str) else param)
            node_outputs.append(output_name)

    def _get_inputs_name(self, node):
        inputs_name = []
        inputs_node = []
        for script_input_ivalue in node.inputs():
            script_input_node = script_input_ivalue.node()
            script_input_unique_id = script_input_ivalue.unique()
            input_name = self.outputs_info[script_input_unique_id]
            inputs_node.append(script_input_node)
            inputs_name.append(input_name)
        return inputs_name, inputs_node

    def data(self, graph, node, uid):
        for output_ivalue in node.outputs():
            script_unique_id = output_ivalue.unique()
            if script_unique_id in self.outputs_info or script_unique_id != uid:
                continue
            node_name = 'x' + str(self.output_index)
            self.outputs_info[script_unique_id] = node_name
            self.output_index += 1
        output_name = self.outputs_info[uid]
        graph.add_layer(
            "fluid.dygraph.base.to_variable",
            inputs={},
            outputs=[node_name],
            value=output_name)
        return [], [output_name]

    def equal(self, graph, node, uid=None, parent_layer=None, index=None):
        if parent_layer is not None and index is not None:
            # block的输出
            input_node_name = self.outputs_info[uid]
            control_output_id = index
            if parent_layer.kernel == "prim.loop":
                control_output_id = index - 1
            output_node_name = parent_layer.outputs[control_output_id]
            graph.add_layer(
                "prim.equal",
                inputs={'input': input_node_name},
                outputs=[output_node_name])
            return [input_node_name], [output_node_name]