pytorch_op_mapper.py 15.9 KB
Newer Older
S
SunAhong1993 已提交
1
# -*- coding:UTF-8 -*-
S
SunAhong1993 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
#   Copyright (c) 2020  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
S
SunAhong1993 已提交
18
from x2paddle.core.util import string
S
SunAhong1993 已提交
19
from x2paddle.core.program import PaddleGraph
S
SunAhong1993 已提交
20 21
from x2paddle.op_mapper.pytorch2paddle import prim
from x2paddle.op_mapper.pytorch2paddle import aten
S
SunAhong1993 已提交
22 23


S
SunAhong1993 已提交
24
class PyTorchOpMapper():
S
SunAhong1993 已提交
25 26
    def __init__(self, decoder):
        self.script = decoder.script
S
SunAhong1993 已提交
27
        self.input_examples = decoder.input_examples
S
SunAhong1993 已提交
28 29 30 31 32
        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
S
SunAhong1993 已提交
33
        self.nn_name2id = {}  # 动态图__init__输出名字中的id,key为kernel类型,value为id
S
SunAhong1993 已提交
34
        self.split_len = {}  # split的长度
S
SunAhong1993 已提交
35 36 37
        self.scope_name_list = list()
        self.scope_name2id = dict()
        self.inputs_info = dict()
S
SunAhong1993 已提交
38
        self.output2id = dict()  # output名字和layer_id的关系,用于lstm去除前面的node
S
SunAhong1993 已提交
39
        # 转换
S
SunAhong1993 已提交
40 41
        if not self.op_checker(decoder.graph):
            raise Exception("Model is not supported yet.")
S
SunAhong1993 已提交
42 43
        self.paddle_graph, _ = self.traverse(decoder.graph)
        self.paddle_graph.set_inputs_info(self.inputs_info)
S
SunAhong1993 已提交
44

S
SunAhong1993 已提交
45
    def op_checker(self, script_graph):
S
SunAhong1993 已提交
46 47 48 49 50
        def _update_op_list(graph):
            for node in graph.nodes():
                op_list.append(node.kind())
                for block in node.blocks():
                    _update_op_list(block)
S
SunAhong1993 已提交
51

S
SunAhong1993 已提交
52 53 54
        op_list = list()
        _update_op_list(script_graph)
        op_list = list(set(op_list))
S
SunAhong1993 已提交
55
        unsupported_ops = []
S
SunAhong1993 已提交
56 57
        for op in op_list:
            func_name = op.replace('::', '_')
W
wjj19950828 已提交
58 59 60 61
            # Processing suffix is "_" situation, eg: aten_relu_ to aten_relu
            # avoid aten::__isnot__ situation
            if func_name[-1] == "_" and func_name[-2] != "_":
                func_name = func_name[:-1]
S
SunAhong1993 已提交
62
            if not (hasattr(prim, func_name) or hasattr(aten, func_name)):
S
SunAhong1993 已提交
63 64 65 66 67
                unsupported_ops.append(op)
        if len(unsupported_ops) == 0:
            return True
        else:
            if len(unsupported_ops) > 0:
S
SunAhong1993 已提交
68 69
                print("\n========= {} OPs are not supported yet ===========".
                      format(len(unsupported_ops)))
S
SunAhong1993 已提交
70 71
            for op in unsupported_ops:
                print("========== {} ============".format(op))
S
SunAhong1993 已提交
72
            return False
S
SunAhong1993 已提交
73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90

    def traverse(self, script_graph, parent_layer=None):
        # 用于获取graph的输入
        def _update_graph_inputs(kind, inputs, outputs):
            # extend只能放更新graph_inputs之前的情况:
            # 1. loop的输出i也是输入;i是输入的原因是:子图中为父图得到的。
            # 2. 在_check_input中需要使用to_variable。
            # extend只能放更新graph_inputs之后的情况:
            # 使用了append。
            if kind != "aten::append":
                current_node_outputs.extend(outputs)
            for name in inputs:
                if name not in current_node_outputs:
                    graph_inputs.append(name)
            if kind == "aten::append":
                current_node_outputs.extend(outputs)

        # 初始化
S
SunAhong1993 已提交
91
        graph = PaddleGraph(source_type="pytorch", parent_layer=parent_layer)
S
SunAhong1993 已提交
92 93
        if "TopLevelTracedModule" in str(type(self.script)):
            graph.set_script(self.script)
S
SunAhong1993 已提交
94 95 96 97
        current_node_outputs = []
        graph_inputs = []
        # 转换输入节点
        if isinstance(script_graph, torch._C.Graph):
S
SunAhong1993 已提交
98
            input_ct = 0
S
SunAhong1993 已提交
99 100
            for i, ivalue in enumerate(script_graph.inputs()):
                node = ivalue.node()
S
SunAhong1993 已提交
101
                if str(ivalue.type()) not in ["Tensor", "Dict[str, Tensor]"]:
S
SunAhong1993 已提交
102 103
                    graph.set_name(str(ivalue.type()).split(".")[-1])
                    continue
S
SunAhong1993 已提交
104 105
                inputs, outputs = self.data(graph, node,
                                            ivalue.unique(), input_ct)
S
SunAhong1993 已提交
106
                input_ct += 1
S
SunAhong1993 已提交
107 108 109 110
        # 转换中间节点
        for node in script_graph.nodes():
            kind = node.kind()
            func_name = kind.replace('::', '_')
W
wjj19950828 已提交
111 112 113 114
            # Processing suffix is "_" situation, eg: aten_relu_ to aten_relu
            # avoid aten::__isnot__ situation
            if func_name[-1] == "_" and func_name[-2] != "_":
                func_name = func_name[:-1]
S
SunAhong1993 已提交
115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147
            if hasattr(prim, func_name):
                func = getattr(prim, func_name)
                inputs, outputs = func(self, graph, node)
                _update_graph_inputs(kind, inputs, outputs)
            elif hasattr(aten, func_name):
                func = getattr(aten, func_name)
                inputs, outputs = func(self, graph, node)
                _update_graph_inputs(kind, 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("equal", inputs, outputs)

        # 设置graph的参数和输出节点
        if isinstance(script_graph, torch._C.Graph):
            graph.set_parameters(self.paddle_params)
            if hasattr(script_graph, 'return_node'):
                inputs_name, inputs_node = self._get_inputs_name(
                    script_graph.return_node())
                graph.outputs = inputs_name
        # 更新split参数
        for layer in graph.layers.values():
S
SunAhong1993 已提交
148
            if layer.kernel == "paddle.split" and "num_or_sections" in layer.attrs \
S
fix  
SunAhong1993 已提交
149
            and not isinstance(layer.attrs["num_or_sections"], int) and len(set(layer.attrs["num_or_sections"])) == 1:
S
SunAhong1993 已提交
150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174
                layer.attrs["num_or_sections"] = self.split_len[layer.outputs[
                    0]]
        return graph, graph_inputs

    def _get_outputs_name(self, node, attr_name=None):
        outputs_name = []
        for output_ivalue in node.outputs():
            script_unique_id = output_ivalue.unique()
            if attr_name is None:
                output_name = 'x' + str(self.output_index)
                if script_unique_id in self.outputs_info:
                    output_name = self.outputs_info[script_unique_id]
            else:
                output_name = attr_name.replace(".", "_")
            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

S
SunAhong1993 已提交
175
    def _check_input(self, graph, node, output_name, node_outputs, scope_name):
S
SunAhong1993 已提交
176 177 178 179
        if node.kind() == "prim::GetAttr":
            param = self.pytorch_params[output_name]
            if isinstance(param, np.ndarray):
                self.paddle_params[output_name] = param
S
SunAhong1993 已提交
180
                layer_id = graph.add_layer(
S
SunAhong1993 已提交
181
                    "self.create_parameter",
S
SunAhong1993 已提交
182 183
                    inputs={},
                    outputs=[output_name],
S
SunAhong1993 已提交
184 185
                    scope_name=scope_name,
                    dtype=string(str(param.dtype)),
S
SunAhong1993 已提交
186 187 188
                    shape=param.shape,
                    default_initializer="paddle.nn.initializer.Constant(value=0.0)"
                )
S
SunAhong1993 已提交
189
                self.output2id[output_name] = layer_id
S
SunAhong1993 已提交
190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206
            else:
                if isinstance(param, dict) and "Tensor" in param and \
                "parent_layer_id" in param:
                    if graph.parent_layer is not None:
                        # 当某个param被2个控制流(if-else)赋值时,else不可以引用if中的赋值结果
                        id1 = param["parent_layer_id"]
                        id2 = graph.parent_layer.id
                        id1_part = id1.split(".")
                        id2_part = id2.split(".")
                        if len(id1_part) >= len(id2_part):
                            for i in range(len(id1_part)):
                                if id1_part[i] == id2_part[i]:
                                    continue
                                else:
                                    if id1_part[i] == "0" and id2_part[
                                            i] == "1":
                                        self.paddle_params[output_name] = param
S
SunAhong1993 已提交
207
                                        layer_id = graph.add_layer(
S
SunAhong1993 已提交
208
                                            "self.create_parameter",
S
SunAhong1993 已提交
209 210
                                            inputs={},
                                            outputs=[output_name],
S
SunAhong1993 已提交
211 212
                                            scope_name=scope_name,
                                            dtype=string(str(param.dtype)),
S
SunAhong1993 已提交
213 214 215
                                            shape=param.shape,
                                            default_initializer="paddle.nn.initializer.Constant(value=0.0)"
                                        )
S
SunAhong1993 已提交
216
                                        node_outputs.append(output_name)
S
SunAhong1993 已提交
217
                                        self.output2id[output_name] = layer_id
S
SunAhong1993 已提交
218 219 220 221 222 223
                                        return
                    # 若if-else外,则可直接引用if-else中的赋值结果
                    graph.add_layer(
                        "prim.constant",
                        inputs={},
                        outputs=[output_name],
S
SunAhong1993 已提交
224
                        scope_name=scope_name,
S
SunAhong1993 已提交
225 226 227 228 229 230
                        value=param["Tensor"])
                else:
                    graph.add_layer(
                        "prim.constant",
                        inputs={},
                        outputs=[output_name],
S
SunAhong1993 已提交
231
                        scope_name=scope_name,
S
SunAhong1993 已提交
232 233 234
                        value=string(param)
                        if isinstance(param, str) else param)
            node_outputs.append(output_name)
S
SunAhong1993 已提交
235 236
        elif node.kind(
        ) == "prim::Constant" and output_name in self.pytorch_params:
S
SunAhong1993 已提交
237 238
            param = self.pytorch_params[output_name]
            self.paddle_params[output_name] = param
S
SunAhong1993 已提交
239
            layer_id = graph.add_layer(
S
SunAhong1993 已提交
240 241 242 243 244
                "self.create_parameter",
                inputs={},
                outputs=[output_name],
                scope_name=scope_name,
                dtype=string(str(param.dtype)),
S
SunAhong1993 已提交
245 246
                shape=param.shape,
                default_initializer="paddle.nn.initializer.Constant(value=0.0)")
S
SunAhong1993 已提交
247
            self.output2id[output_name] = layer_id
S
SunAhong1993 已提交
248 249 250 251 252 253 254 255 256 257 258 259

    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

S
SunAhong1993 已提交
260 261
    def data(self, graph, node, uid, input_ct):
        scope_name = self.normalize_scope_name(node)
S
SunAhong1993 已提交
262 263 264 265 266 267 268 269 270
        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(
S
SunAhong1993 已提交
271
            "paddle.to_tensor",
S
SunAhong1993 已提交
272 273
            inputs={},
            outputs=[node_name],
S
SunAhong1993 已提交
274 275 276
            scope_name=scope_name,
            data=output_name)
        if self.input_examples is not None:
C
channingss 已提交
277
            input_np = self.input_examples[input_ct].cpu().detach().numpy()
S
SunAhong1993 已提交
278 279
            self.inputs_info[
                output_name] = [list(input_np.shape), str(input_np.dtype)]
S
SunAhong1993 已提交
280 281 282
        return [], [output_name]

    def equal(self, graph, node, uid=None, parent_layer=None, index=None):
S
SunAhong1993 已提交
283
        scope_name = self.normalize_scope_name(node)
S
SunAhong1993 已提交
284 285 286 287 288 289 290 291
        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]
            current_outputs = [output_node_name]
S
SunAhong1993 已提交
292 293
            self._check_input(graph, node, input_node_name, current_outputs,
                              scope_name)
S
SunAhong1993 已提交
294 295 296
            graph.add_layer(
                "prim.equal",
                inputs={'input': input_node_name},
S
SunAhong1993 已提交
297 298
                outputs=[output_node_name],
                scope_name=scope_name)
S
SunAhong1993 已提交
299
            return [input_node_name], current_outputs
S
SunAhong1993 已提交
300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324

    def normalize_scope_name(self, node):
        """ 对scope的名字进行标准化。
        """
        scope_name = node.scopeName()
        if scope_name == "":
            return scope_name
        scope_name_part = scope_name.split("/")
        for index in range(len(scope_name_part) - 1):
            if scope_name_part[index] in scope_name_part[index + 1]:
                continue
            last_name_segments = scope_name_part[index].split(".")
            name_segments = scope_name_part[index + 1].split(".")
            for j, name in enumerate(last_name_segments):
                name_segments[j] = name
            scope_name_part[index + 1] = ".".join(name_segments)
        last_name = scope_name_part[-1]
        name_segments = last_name.split(".")
        for i, ns in enumerate(name_segments):
            if i not in self.scope_name2id:
                self.scope_name2id[i] = dict()
            if ns not in self.scope_name2id[i]:
                self.scope_name2id[i][ns] = 0
        real_scope_name = "/".join(name_segments[1:])
        real_father_scope_name = "/".join(name_segments[1:-1])
S
SunAhong1993 已提交
325

S
SunAhong1993 已提交
326 327 328 329 330 331
        for i, ns in enumerate(name_segments):
            if i == 0:
                continue
            if self.scope_name2id[i][ns] != 0:
                name_segments[i] = name_segments[i] + \
                "__{}".format(self.scope_name2id[i][ns])
S
SunAhong1993 已提交
332
            prefix_scope_name = "/".join(name_segments[1:i + 1])
S
SunAhong1993 已提交
333 334
            is_found = False
            for j in range(len(self.scope_name_list)):
S
SunAhong1993 已提交
335
                last_scope_name = self.scope_name_list[-1 - j]
S
SunAhong1993 已提交
336 337
                if last_scope_name.startswith(prefix_scope_name + "/") \
                        or last_scope_name == prefix_scope_name:
S
SunAhong1993 已提交
338
                    if j != 0:  # and i != len(name_segments) - 1:
S
SunAhong1993 已提交
339 340 341 342 343 344 345 346 347 348 349
                        is_found = True
                        origin_name_segment_i = name_segments[i].split("__")[0]
                        self.scope_name2id[i][origin_name_segment_i] += 1
                        name_segments[i] = origin_name_segment_i + \
                            "__" + str(self.scope_name2id[i][origin_name_segment_i])
                    break
            if is_found:
                break
        real_scope_name = "/".join(name_segments[1:])
        self.scope_name_list.append(real_scope_name)
        return real_scope_name