program.py 23.2 KB
Newer Older
S
SunAhong1993 已提交
1
# -*- coding:UTF-8 -*-
S
SunAhong1993 已提交
2
#   Copyright (c) 2020  PaddlePaddle Authors. All Rights Reserved.
J
jiangjiajun 已提交
3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
#
# 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.

from __future__ import print_function
from __future__ import division
S
SunAhong1993 已提交
18
import paddle
C
Channingss 已提交
19
import collections
S
SunAhong1993 已提交
20
import sys
J
jiangjiajun 已提交
21
import os
S
SunAhong1993 已提交
22 23
import six
import pickle
S
SunAhong1993 已提交
24
from os import path as osp
S
SunAhong1993 已提交
25
from x2paddle.core.util import *
J
jiangjiajun 已提交
26 27 28


class PaddleLayer(object):
S
SunAhong1993 已提交
29
    def __init__(self, id, kernel, inputs, outputs, scope_name="", **kwargs):
J
jiangjiajun 已提交
30 31 32 33 34
        assert isinstance(
            inputs,
            dict), "parameter 'inputs' for PaddleLayer should be type of dict"
        assert isinstance(
            outputs,
S
SunAhong1993 已提交
35 36
            list), "parameter 'outputs' for PaddleLayer should be type of list"
        for k, v in inputs.items():
C
Channingss 已提交
37
            if isinstance(v, (list, tuple)):
S
SunAhong1993 已提交
38 39 40 41 42 43 44 45 46 47 48 49
                for i in v:
                    assert isinstance(
                        i, six.string_types
                    ), "value in inputs should be type of string or list of string"
            else:
                assert isinstance(v, six.string_types) or isinstance(
                    v, list
                ), "value in inputs should be type of string or list of string"
        for v in outputs:
            assert isinstance(
                v, six.
                string_types), "elements in outputs should be type of string"
J
jiangjiajun 已提交
50 51 52
        self.kernel = kernel
        self.inputs = inputs
        self.outputs = outputs
S
SunAhong1993 已提交
53
        self.scope_name = scope_name
J
jiangjiajun 已提交
54
        self.attrs = kwargs
S
SunAhong1993 已提交
55 56 57 58 59
        self.id = id
        self.blocks = list()

    def add_block(self, block):
        self.blocks.append(block)
J
jiangjiajun 已提交
60 61


S
SunAhong1993 已提交
62
class PaddleGraph(object):
S
SunAhong1993 已提交
63
    def __init__(self, source_type=None, parent_layer=None):
C
Channingss 已提交
64
        self.layers = collections.OrderedDict()
J
jiangjiajun 已提交
65 66 67 68 69
        self.edges_out = dict()
        self.edges_in = dict()
        self.inputs = list()
        self.outputs = list()
        self.parameters = dict()
S
SunAhong1993 已提交
70
        self.parent_layer = parent_layer
S
SunAhong1993 已提交
71 72
        self.source_type = source_type
        self.custom_code = None
S
SunAhong1993 已提交
73
        self.inputs_info = None
S
SunAhong1993 已提交
74
        self.has_unpack = False
S
SunAhong1993 已提交
75

S
SunAhong1993 已提交
76
    def set_name(self, name):
S
SunAhong1993 已提交
77
        self.name = name.replace("-", "_").replace("/", "_")
S
SunAhong1993 已提交
78 79 80

    def set_parameters(self, parameters):
        self.parameters = parameters
S
SunAhong1993 已提交
81

S
SunAhong1993 已提交
82 83
    def set_custom(self, custom_code):
        self.custom_code = custom_code
S
SunAhong1993 已提交
84

S
SunAhong1993 已提交
85 86
    def set_inputs_info(self, inputs_info):
        self.inputs_info = inputs_info
S
SunAhong1993 已提交
87

S
SunAhong1993 已提交
88 89
    def set_script(self, script):
        self.script = script
S
SunAhong1993 已提交
90 91

    def clear(self):
C
Channingss 已提交
92
        self.layers = collections.OrderedDict()
S
SunAhong1993 已提交
93 94 95 96 97 98 99 100 101
        self.edges_out = dict()
        self.edges_in = dict()
        self.inputs = list()
        self.outputs = list()
        self.parameters = dict()

    def clear_edges(self):
        self.edges_out = dict()
        self.edges_in = dict()
J
jiangjiajun 已提交
102

S
SunAhong1993 已提交
103
    def add_layer(self, kernel, inputs, outputs, scope_name="", **kwargs):
S
SunAhong1993 已提交
104 105 106 107 108
        layer_id = str(len(self.layers))
        if self.parent_layer is not None:
            layer_id = "{}.{}.{}".format(self.parent_layer.id,
                                         len(self.parent_layer.blocks),
                                         layer_id)
S
SunAhong1993 已提交
109 110
        layer = PaddleLayer(
            layer_id, kernel, inputs, outputs, scope_name=scope_name, **kwargs)
S
SunAhong1993 已提交
111
        self.layers[layer_id] = layer
112
        if layer.kernel in ["prim.list_unpack", "prim.tuple_unpack"]:
S
SunAhong1993 已提交
113
            self.has_unpack = True
S
SunAhong1993 已提交
114
        return layer_id
J
jiangjiajun 已提交
115

J
jiangjiajun 已提交
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 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
    def del_layer(self, layer_id):
        layer = self.layers[layer_id]
        outputs = self.edges_out.get(layer_id, [])
        inputs = self.edges_in.get(layer_id, [])

        assert len(
            inputs) <= 1, "There should be 0 or 1 input for deleted layer."

        if len(inputs) == 0:
            for out in outputs:
                while layer_id in self.edges_in[out]:
                    index = self.edges_in[out].index(layer_id)
                    del self.edges_in[out][index]

                input_keys = list(self.layers[out].inputs.keys())
                for k in input_keys:
                    if self.layers[out].inputs[k] == layer.outputs[0]:
                        del self.layers[out].inputs[k]

            del self.layers[layer_id]
            if layer_id in self.edges_in:
                del self.edges_in[layer_id]
            if layer_id in self.edges_out:
                del self.edges_out[layer_id]
            return

        # 将所有输出layer的输入layer进行替换
        for out in outputs:
            for i in range(len(self.edges_in[out])):
                if self.edges_in[out][i] == layer_id:
                    self.edges_in[out][i] = inputs[0]

        # 将输出layer赋给输入layer的输出
        replace_index = self.edges_out[inputs[0]].index(layer_id)
        del self.edges_out[inputs[0]][replace_index]
        for i, out in enumerate(outputs):
            self.edges_out[inputs[0]].insert(replace_index + i, out)
            for k, v in self.layers[out].inputs.items():
                if v == layer.outputs[0]:
                    self.layers[out].inputs[k] = list(layer.inputs.values())[0]

        del self.layers[layer_id]
        if layer_id in self.edges_out:
            del self.edges_out[layer_id]
        if layer_id in self.edges_in:
            del self.edges_in[layer_id]

S
SunAhong1993 已提交
163 164 165 166 167 168
    def build(self, inputs=None, outputs=None):
        self.clear_edges()
        outputs_from_nodes = dict()
        for layer_id, layer in self.layers.items():
            for input_key, input_var in layer.inputs.items():
                vs = input_var
C
Channingss 已提交
169
                if not isinstance(vs, (list, tuple)):
S
SunAhong1993 已提交
170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188
                    vs = [vs]
                for v in vs:
                    assert v in outputs_from_nodes or (
                        inputs is not None and v in list(inputs.values())
                    ) or (
                        outputs is not None and v in outputs
                    ), "Couldn't find {} in previous layers, the layers should be make by topological sort".format(
                        v)
                    if v in outputs_from_nodes:
                        in_layer_id = outputs_from_nodes[v]
                    else:
                        in_layer_id = -1
                    if in_layer_id not in self.edges_out:
                        self.edges_out[in_layer_id] = list()
                    self.edges_out[in_layer_id].append(layer_id)

                    if layer_id not in self.edges_in:
                        self.edges_in[layer_id] = list()
                    self.edges_in[layer_id].append(in_layer_id)
J
jiangjiajun 已提交
189
            for output in layer.outputs:
S
SunAhong1993 已提交
190
                outputs_from_nodes[output] = layer_id
J
jiangjiajun 已提交
191

S
SunAhong1993 已提交
192 193 194 195 196 197
            # 将block的输出用于父图
            if inputs is not None and outputs is not None and set(
                    layer.outputs).issubset(outputs):
                if layer_id not in self.edges_out:
                    self.edges_out[layer_id] = list()
                self.edges_out[layer_id].append(-1)
J
jiangjiajun 已提交
198

S
SunAhong1993 已提交
199 200 201 202
            # 处理子图
            if len(layer.blocks) > 0:
                for block in layer.blocks:
                    block.build(layer.inputs, layer.outputs)
J
jiangjiajun 已提交
203

S
SunAhong1993 已提交
204 205 206 207 208 209 210
        # 删除不必要的节点
        invalid_list = list()
        for layer_id, layer in self.layers.items():
            if len(self.layers) > 1:
                if self.edges_in.get(layer_id, 0) == 0 and self.edges_out.get(
                        layer_id, 0) == 0 and layer.kernel != "prim.assert" \
                        and layer.kernel != "prim.exception" \
S
SunAhong1993 已提交
211 212
                        and layer.kernel != "prim.warnings" \
                        and layer.outputs[0] not in self.outputs:
S
SunAhong1993 已提交
213 214
                    if layer.kernel == "paddle.to_tensor" and layer.outputs[
                            0] in self.inputs_info:
S
SunAhong1993 已提交
215
                        self.inputs_info.pop(layer.outputs[0])
S
SunAhong1993 已提交
216 217
                    if layer.outputs[0] in self.inputs:
                        self.inputs.pop(self.inputs.index(layer.outputs[0]))
S
SunAhong1993 已提交
218 219 220
                    invalid_list.append(layer_id)
        for layer_id in invalid_list:
            self.layers.pop(layer_id)
J
jiangjiajun 已提交
221

S
SunAhong1993 已提交
222
        self.get_inputs()
S
SunAhong1993 已提交
223

S
SunAhong1993 已提交
224 225
        if len(self.outputs) == 0:
            self.get_outputs()
J
jiangjiajun 已提交
226

S
SunAhong1993 已提交
227 228 229 230 231 232 233 234 235 236
    def get_global_layers(self):
        # 该全局layers的信息是按照拓扑排序组成的
        def update(layers):
            global_layers = dict()
            for layer_id, layer in layers.items():
                global_layers[layer_id] = layer
                for block in layer.blocks:
                    block_global_layers = update(block.layers)
                    global_layers.update(block_global_layers)
            return global_layers
S
SunAhong1993 已提交
237

S
SunAhong1993 已提交
238
        return update(self.layers)
S
SunAhong1993 已提交
239

240
    def gen_model(self, save_dir, jit_type=None, enable_code_optim=True):
S
SunAhong1993 已提交
241
        if not osp.exists(save_dir):
S
SunAhong1993 已提交
242
            os.makedirs(save_dir)
S
SunAhong1993 已提交
243
        if jit_type == "trace":
244
            if not self.has_unpack and enable_code_optim:
S
SunAhong1993 已提交
245 246 247 248 249
                from x2paddle.optimizer.pytorch_code_optimizer import HierarchicalTree
                hierarchical_tree = HierarchicalTree(self)
                for layer_id, layer in self.layers.items():
                    hierarchical_tree.insert(layer)
                hierarchical_tree.save_source_files(save_dir)
S
SunAhong1993 已提交
250
                self.dump_parameter(save_dir)
S
SunAhong1993 已提交
251
            else:
S
SunAhong1993 已提交
252 253
                self.gen_code(save_dir)
                self.dump_parameter(save_dir)
S
SunAhong1993 已提交
254
        else:
255
            if self.source_type == "pytorch" and enable_code_optim:
S
SunAhong1993 已提交
256
                from x2paddle.optimizer.pytorch_code_optimizer import ModuleGraph
S
SunAhong1993 已提交
257 258
                module_graph = ModuleGraph(self)
                module_graph.save_source_files(save_dir)
S
SunAhong1993 已提交
259
                self.dump_parameter(save_dir)
S
SunAhong1993 已提交
260
            else:
S
SunAhong1993 已提交
261 262
                self.gen_code(save_dir)
                self.dump_parameter(save_dir)
S
SunAhong1993 已提交
263
        # 动转静
S
SunAhong1993 已提交
264
        code_path = osp.join(osp.abspath(save_dir), "x2paddle_code.py")
S
SunAhong1993 已提交
265 266
        print("Exporting inference model from python code ('{}')... \n".format(
            code_path))
S
SunAhong1993 已提交
267
        if len(self.inputs_info) > 0:
S
SunAhong1993 已提交
268 269 270 271 272
            input_shapes = list()
            input_types = list()
            for input_name in self.inputs:
                input_shapes.append(self.inputs_info[input_name][0])
                input_types.append(self.inputs_info[input_name][1])
S
SunAhong1993 已提交
273 274
            try:
                self.dygraph2static(save_dir, input_shapes, input_types)
S
SunAhong1993 已提交
275
            except Exception as e:
S
SunAhong1993 已提交
276 277 278
                print(
                    "Fail to generate inference model! Problem happend while export inference model from python code '{}';\n".
                    format(code_path))
S
SunAhong1993 已提交
279 280
                print("===================Error Information===============")
                raise e
S
SunAhong1993 已提交
281

S
SunAhong1993 已提交
282
    def get_inputs(self):
S
SunAhong1993 已提交
283 284 285 286 287
        def update(layers):
            for layer_id, layer in layers.items():
                if self.edges_in.get(layer_id, 0) == 0 and self.edges_out.get(
                        layer_id, 0) == 0:
                    continue
S
SunAhong1993 已提交
288 289
                if layer.kernel == "paddle.to_tensor":
                    data = layer.attrs["data"]
S
SunAhong1993 已提交
290
                    self.inputs.append(data)
S
SunAhong1993 已提交
291 292
                if len(layer.blocks) > 0:
                    for block in layer.blocks:
S
SunAhong1993 已提交
293
                        block.get_inputs()
S
SunAhong1993 已提交
294 295 296
                        self.inputs.extend(block.inputs)

        update(self.layers)
S
SunAhong1993 已提交
297 298 299 300 301 302 303
        new_inputs = list()
        for input_name in self.inputs:
            if input_name in new_inputs:
                continue
            new_inputs.append(input_name)
        self.inputs = new_inputs
        if self.source_type == "pytorch" and self.inputs is not None:
S
SunAhong1993 已提交
304 305
            self.inputs.sort()

S
SunAhong1993 已提交
306
    def get_outputs(self):
S
SunAhong1993 已提交
307 308 309 310 311
        for layer_id, layer in self.layers.items():
            if self.edges_in.get(layer_id, 0) == 0 and self.edges_out.get(
                    layer_id, 0) == 0:
                continue
            if self.edges_out.get(layer_id, 0) == 0:
S
SunAhong1993 已提交
312

S
SunAhong1993 已提交
313
                for i, output_name in enumerate(layer.outputs):
S
SunAhong1993 已提交
314 315
                    if ("paddle.nn" in layer.kernel and
                            "functional" not in layer.kernel):
S
SunAhong1993 已提交
316 317 318 319
                        if i == 0:
                            continue
                    if output_name not in self.outputs:
                        self.outputs.append(output_name)
S
SunAhong1993 已提交
320

S
SunAhong1993 已提交
321
    def gen_code(self, code_dir=None, indent=2):
S
SunAhong1993 已提交
322 323 324 325 326 327 328 329 330
        # 去除to_tensor的layer
        invalid_list = list()
        for layer_id, layer in self.layers.items():
            if layer.kernel == "paddle.to_tensor":
                if layer.attrs["data"] in self.inputs:
                    invalid_list.append(layer_id)
        for layer_id in invalid_list:
            self.layers.pop(layer_id)

S
SunAhong1993 已提交
331 332 333 334 335 336 337 338 339 340 341
        def gen_codes(code_list, indent=0):
            indent_blank = "    " * indent
            codes = []
            for code_line in code_list:
                if code_line.strip() == "":
                    codes.append('\n')
                else:
                    codes.append(indent_blank + code_line + '\n')
            return codes

        def gen_head():
S
SunAhong1993 已提交
342
            if self.source_type == "caffe":
S
SunAhong1993 已提交
343
                custom_import = "from x2paddle.op_mapper.caffe2paddle " + \
S
SunAhong1993 已提交
344
                                 "import caffe_custom_layer as x2paddle_nn"
S
SunAhong1993 已提交
345
            elif self.source_type == "pytorch":
S
SunAhong1993 已提交
346
                custom_import = "from x2paddle.op_mapper.pytorch2paddle " + \
S
SunAhong1993 已提交
347
                                 "import pytorch_custom_layer as x2paddle_nn"
S
SunAhong1993 已提交
348
            elif self.source_type == "onnx":
S
SunAhong1993 已提交
349
                custom_import = "from x2paddle.op_mapper.onnx2paddle " + \
S
SunAhong1993 已提交
350
                                 "import onnx_custom_layer as x2paddle_nn"
S
SunAhong1993 已提交
351 352
            else:
                custom_import = ""
S
SunAhong1993 已提交
353 354 355
            self.head = gen_codes(
                [
                    "import paddle",
S
SunAhong1993 已提交
356
                    "import math",
S
SunAhong1993 已提交
357
                    custom_import,
S
SunAhong1993 已提交
358
                    "",
S
SunAhong1993 已提交
359
                    "class {}(paddle.nn.Layer):".format(self.name),
S
SunAhong1993 已提交
360 361 362
                ],
                indent=0)
            input_data_name = ', '.join(self.inputs)
S
SunAhong1993 已提交
363
            self.init_func.extend(gen_codes(["def __init__(self):"], indent=1))
S
SunAhong1993 已提交
364 365 366 367 368 369 370
            self.init_func.extend(
                gen_codes(
                    ["super({}, self).__init__()".format(self.name)], indent=2))
            self.forward_func.extend(
                gen_codes(
                    ["def forward(self, {}):".format(input_data_name)],
                    indent=1))
S
SunAhong1993 已提交
371

S
SunAhong1993 已提交
372
        def gen_main_code(code_dir):
S
SunAhong1993 已提交
373 374 375 376
            input_data_name = ', '.join(self.inputs)
            self.run_func = gen_codes(
                [
                    "",
S
SunAhong1993 已提交
377
                    "def main({}):".format(input_data_name),
S
SunAhong1993 已提交
378
                ], indent=0)
S
SunAhong1993 已提交
379
            comment_list = list()
S
SunAhong1993 已提交
380 381
            comment_list.append("# There are {} inputs.".format(
                len(self.inputs_info)))
S
SunAhong1993 已提交
382
            for k, v in self.inputs_info.items():
S
SunAhong1993 已提交
383 384 385 386
                comment_list.append("# {}: shape-{}, type-{}.".format(k, v[0],
                                                                      v[1]))
            self.run_func.extend(gen_codes(comment_list, indent=1))
            use_structured_name = False if self.source_type in ["tf"] else True
S
SunAhong1993 已提交
387 388
            self.run_func.extend(
                gen_codes(
S
SunAhong1993 已提交
389 390
                    [
                        "paddle.disable_static()",
391
                        "params = paddle.load(r'{}')".format(
S
SunAhong1993 已提交
392 393 394 395 396 397
                            osp.join(osp.abspath(code_dir), "model.pdparams")),
                        "model = {}()".format(self.name),
                        "model.set_dict(params, use_structured_name={})".format(
                            use_structured_name), "model.eval()",
                        "out = model({})".format(input_data_name), "return out"
                    ],
S
SunAhong1993 已提交
398
                    indent=1))
S
SunAhong1993 已提交
399 400

        def write_code(code_dir):
S
SunAhong1993 已提交
401
            f = open(osp.join(code_dir, 'x2paddle_code.py'), 'w')
S
SunAhong1993 已提交
402 403 404 405 406 407 408 409 410 411 412 413
            for code_line in self.head:
                f.write(code_line)
            init_writen_codes = []
            for code_line in self.init_func:
                if code_line in init_writen_codes:
                    continue
                f.write(code_line)
                init_writen_codes.append(code_line)
            f.write("\n")
            return_code = "return {}".format(", ".join(self.outputs))
            self.forward_func.extend(gen_codes([return_code], indent=2))
            for code_line in self.forward_func:
S
SunAhong1993 已提交
414 415
                if "assert [1, 1] == 1 or [1, 1] == [1, 1], 'The [1, 1] must be [1, [1, 1]]!'" in code_line:
                    continue
S
SunAhong1993 已提交
416
                f.write(code_line)
S
SunAhong1993 已提交
417 418
            for code_line in self.run_func:
                f.write(code_line)
S
SunAhong1993 已提交
419 420 421 422 423 424 425 426
            f.close()

        self.init_func = []
        self.forward_func = []
        if indent == 2 and code_dir is not None:
            gen_head()

        for layer_id, layer in self.layers.items():
S
SunAhong1993 已提交
427 428
            if layer.kernel.startswith("paddle"):
                remove_default_attrs(layer.kernel, layer.attrs)
S
SunAhong1993 已提交
429
            if ("paddle.nn" in layer.kernel and "functional" not in layer.kernel
S
SunAhong1993 已提交
430
                ) or layer.kernel == "paddle.to_tensor" or \
S
SunAhong1993 已提交
431
                layer.kernel.startswith("custom_layer"):
S
SunAhong1993 已提交
432 433
                line = "{}".format(
                    layer.outputs[0]
S
SunAhong1993 已提交
434 435
                ) if layer.kernel == "paddle.to_tensor" and not layer.attrs[
                    "data"].startswith("params[") else "self.{}".format(
S
SunAhong1993 已提交
436
                        layer.outputs[0])
S
SunAhong1993 已提交
437
                if layer.kernel.startswith("custom_layer"):
S
SunAhong1993 已提交
438 439
                    line += "= x2paddle_nn.{}(".format(
                        layer.kernel.split(":")[-1])
S
SunAhong1993 已提交
440 441
                else:
                    line += " = {}(".format(layer.kernel)
S
SunAhong1993 已提交
442 443 444 445 446
                for k, v in layer.attrs.items():
                    line += "{}={}, ".format(k, v)
                line = line.strip(", ")
                line += ")"

S
SunAhong1993 已提交
447 448
                if layer.kernel == "paddle.to_tensor" and not layer.attrs[
                        "data"].startswith("params["):
S
SunAhong1993 已提交
449 450 451 452 453 454 455 456 457 458
                    self.forward_func.extend(gen_codes([line], indent=indent))
                    continue
                else:
                    self.init_func.extend(gen_codes([line], indent=2))

                if len(layer.outputs) == 1:
                    line = layer.outputs[0]
                elif len(layer.outputs) == 2:
                    line = layer.outputs[1]
                else:
C
Channingss 已提交
459
                    if layer.kernel in ["paddle.nn.LSTM"]:
S
SunAhong1993 已提交
460 461
                        line = "{}, ({})".format(layer.outputs[1],
                                                 ', '.join(layer.outputs[-2:]))
S
SunAhong1993 已提交
462 463
                    else:
                        line = ','.join(layer.outputs[1:])
S
SunAhong1993 已提交
464 465
                if layer.kernel == "paddle.to_tensor" and layer.attrs[
                        "data"].startswith("params["):
S
SunAhong1993 已提交
466 467 468
                    line += " = self.{}".format(layer.outputs[0])
                else:
                    line += " = self.{}(".format(layer.outputs[0])
469 470 471 472 473 474 475
                    for v in layer.inputs.values():
                        if isinstance(v, list):
                            line += "[{}], ".format(", ".join(v))
                        elif isinstance(v, tuple):
                            line += "({}), ".format(", ".join(v))
                        else:
                            line += "{}, ".format(v)
S
SunAhong1993 已提交
476 477
                    line = line.strip(", ")
                    line += ")"
S
SunAhong1993 已提交
478
                self.forward_func.extend(gen_codes([line], indent=indent))
S
SunAhong1993 已提交
479 480
            elif "prim" in layer.kernel:
                func_name = layer.kernel.replace(".", "_")
S
SunAhong1993 已提交
481
                from x2paddle.op_mapper.pytorch2paddle import prim2code
S
SunAhong1993 已提交
482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499
                if hasattr(prim2code, func_name):
                    func = getattr(prim2code, func_name)
                    func(
                        layer,
                        indent=indent,
                        init_func=self.init_func,
                        forward_func=self.forward_func)
                else:
                    raise Exception(
                        "The kind {} in paddle model is not supported yet.".
                        format(layer.kernel))
            else:
                if len(layer.outputs) == 1:
                    line = layer.outputs[0]
                else:
                    line = ','.join(layer.outputs)
                line += " = {}(".format(layer.kernel)
                for k, v in layer.inputs.items():
S
SunAhong1993 已提交
500 501
                    if isinstance(v, list):
                        line += "{}=[{}], ".format(k, ", ".join(v))
C
Channingss 已提交
502 503
                    elif isinstance(v, tuple):
                        line += "{}=({}), ".format(k, ", ".join(v))
S
SunAhong1993 已提交
504
                    else:
S
SunAhong1993 已提交
505 506 507 508
                        if k == "args":
                            line += v
                        else:
                            line += "{}={}, ".format(k, v)
S
SunAhong1993 已提交
509 510 511 512
                for k, v in layer.attrs.items():
                    line += "{}={}, ".format(k, v)
                line = line.strip(", ")
                line += ")"
S
SunAhong1993 已提交
513
                if layer.kernel == "self.create_parameter":
S
SunAhong1993 已提交
514
                    self.init_func.extend(gen_codes(["self." + line], indent=2))
S
SunAhong1993 已提交
515 516 517 518 519 520 521
                    self.forward_func.extend(
                        gen_codes(
                            [
                                "{} = self.{}".format(layer.outputs[0],
                                                      layer.outputs[0])
                            ],
                            indent=indent))
S
SunAhong1993 已提交
522 523
                else:
                    self.forward_func.extend(gen_codes([line], indent=indent))
S
SunAhong1993 已提交
524
        if indent == 2 and code_dir is not None:
S
SunAhong1993 已提交
525
            gen_main_code(code_dir)
S
SunAhong1993 已提交
526 527 528 529
            write_code(code_dir)
        else:
            return self.init_func, self.forward_func

S
SunAhong1993 已提交
530
    def dump_parameter(self, code_dir):
S
SunAhong1993 已提交
531
        save_path = osp.join(code_dir, 'model.pdparams')
S
SunAhong1993 已提交
532
        paddle.save(self.parameters, save_path)
J
jiangjiajun 已提交
533

S
SunAhong1993 已提交
534
    def dygraph2static(self, save_dir, input_shapes=[], input_types=[]):
S
SunAhong1993 已提交
535 536 537 538
        sepc_list = list()
        for i, name in enumerate(self.inputs):
            sepc_list.append(
                paddle.static.InputSpec(
S
SunAhong1993 已提交
539
                    shape=input_shapes[i], name=name, dtype=input_types[i]))
S
SunAhong1993 已提交
540 541 542
        path = osp.abspath(save_dir)
        sys.path.insert(0, save_dir)
        import x2paddle_code
S
SunAhong1993 已提交
543
        paddle.disable_static()
S
SunAhong1993 已提交
544
        restore = paddle.load(osp.join(save_dir, "model.pdparams"))
S
SunAhong1993 已提交
545
        model = getattr(x2paddle_code, self.name)()
C
Channingss 已提交
546
        if self.source_type in ["tf"]:
S
SunAhong1993 已提交
547 548 549
            model.set_dict(restore, use_structured_name=False)
        else:
            model.set_dict(restore)
S
SunAhong1993 已提交
550 551
        model.eval()
        static_model = paddle.jit.to_static(model, input_spec=sepc_list)
S
SunAhong1993 已提交
552
        try:
S
SunAhong1993 已提交
553 554
            paddle.jit.save(static_model,
                            osp.join(save_dir, "inference_model/model"))
S
SunAhong1993 已提交
555 556
        except ValueError as e:
            if str(e) == "'target_vars' should be a list of Variable.":
S
SunAhong1993 已提交
557 558 559
                print(
                    "[DyGraph2StaticGraph Error] Can not convert the dygraph to static! The output of PyTorch mustbe Variable or a list of Variable."
                )
S
SunAhong1993 已提交
560
            else:
S
SunAhong1993 已提交
561
                print(e)
S
fix  
SunAhong1993 已提交
562
                exit(0)