op_mapper.py 6.1 KB
Newer Older
J
jiangjiajun 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
#   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.
J
modify  
jiangjiajun 已提交
14
from paddle.fluid.proto import framework_pb2
J
jiangjiajun 已提交
15
from x2paddle.core.util import *
J
jiangjiajun 已提交
16
import inspect
J
jiangjiajun 已提交
17 18 19
import os


J
jiangjiajun 已提交
20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48
def export_paddle_param(param, param_name, dir):
    dtype_map = {
        "int16": [framework_pb2.VarType.INT16, 'h'],
        "int32": [framework_pb2.VarType.INT32, 'i'],
        "int64": [framework_pb2.VarType.INT64, 'q'],
        "float16": [framework_pb2.VarType.FP16, 'e'],
        "float32": [framework_pb2.VarType.FP32, 'f'],
        "float64": [framework_pb2.VarType.FP64, 'd']
    }
    shape = param.shape
    if len(shape) == 0:
        assert param.size == 1, "Unexpected situation happend!"
        shape = [1]
    assert str(param.dtype) in dtype_map, "Unknown dtype of params."

    fp = open(os.path.join(dir, param_name), 'wb')
    numpy.array([0], dtype='int32').tofile(fp)
    numpy.array([0], dtype='int64').tofile(fp)
    numpy.array([0], dtype='int32').tofile(fp)
    tensor_desc = framework_pb2.VarType.TensorDesc()
    tensor_desc.data_type = dtype_map[str(param.dtype)][0]
    tensor_desc.dims.extend(shape)
    desc_size = tensor_desc.ByteSize()
    numpy.array([desc_size], dtype='int32').tofile(fp)
    fp.write(tensor_desc.SerializeToString())
    param.tofile(fp)
    fp.close()


J
jiangjiajun 已提交
49 50 51 52 53 54
class OpMapper(object):
    def __init__(self):
        self.paddle_codes = ""
        self.tab = "    "
        self.net_code = list()
        self.weights = dict()
J
jiangjiajun 已提交
55 56
        self.inputs = list()
        self.outputs = list()
J
jiangjiajun 已提交
57

J
jiangjiajun 已提交
58 59 60 61 62 63 64 65 66 67
    def op_checker(self):
        unsupported_ops = set()
        for node_name in self.graph.topo_sort:
            node = self.graph.get_node(node_name)
            op = node.layer_type
            if not hasattr(self, op):
                unsupported_ops.add(op)
        if len(unsupported_ops) == 0:
            return True
        else:
J
jiangjiajun 已提交
68 69
            print("There are {} ops not supported yet, list as below".format(
                len(unsupported_ops)))
J
jiangjiajun 已提交
70 71 72 73
            for op in unsupported_ops:
                print(op)
            return False

J
jiangjiajun 已提交
74 75 76
    def add_codes(self, codes, indent=0):
        if isinstance(codes, list):
            for code in codes:
J
jiangjiajun 已提交
77 78
                self.paddle_codes += (self.tab * indent + code.strip('\n') +
                                      '\n')
J
jiangjiajun 已提交
79
        elif isinstance(codes, str):
J
jiangjiajun 已提交
80
            self.paddle_codes += (self.tab * indent + codes.strip('\n') + '\n')
J
jiangjiajun 已提交
81 82 83 84 85 86 87 88 89
        else:
            raise Exception("Unknown type of codes")

    def add_heads(self):
        self.add_codes("from paddle.fluid.initializer import Constant")
        self.add_codes("from paddle.fluid.param_attr import ParamAttr")
        self.add_codes("import paddle.fluid as fluid")
        self.add_codes("")

J
jiangjiajun 已提交
90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105
    def save_inference_model(self, save_dir):
        self.save_python_model(save_dir)

        import sys
        import paddle.fluid as fluid
        py_code_dir = os.path.join(save_dir, "model_with_code")
        sys.path.append(py_code_dir)
        import model
        try:
            inputs, outputs = model.x2paddle_net()
            input_names = [input.name for input in inputs]
            exe = fluid.Executor(fluid.CPUPlace())
            exe.run(fluid.default_startup_program())

            def if_exist(var):
                b = os.path.exists(
J
jiangjiajun 已提交
106
                    os.path.join(os.path.join(py_code_dir, var.name)))
J
jiangjiajun 已提交
107 108 109
                return b

            fluid.io.load_vars(exe,
J
jiangjiajun 已提交
110
                               py_code_dir,
J
jiangjiajun 已提交
111 112 113 114 115 116 117 118
                               fluid.default_main_program(),
                               predicate=if_exist)

            fluid.io.save_inference_model(dirname=os.path.join(
                save_dir, "inference_model"),
                                          feeded_var_names=input_names,
                                          target_vars=outputs,
                                          executor=exe,
J
jiangjiajun 已提交
119
                                          params_filename=None)
J
jiangjiajun 已提交
120 121 122 123
        except:
            raise Exception(
                "Paddle code was saved in {}/model.py, but seems there's wrong exist, please check model.py manually."
                .format(py_code_dir))
J
jiangjiajun 已提交
124 125

    def save_python_model(self, save_dir):
J
jiangjiajun 已提交
126 127 128 129 130 131 132
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        py_code_dir = os.path.join(save_dir, "model_with_code")
        if not os.path.exists(py_code_dir):
            os.makedirs(py_code_dir)

J
jiangjiajun 已提交
133
        for name, param in self.weights.items():
J
jiangjiajun 已提交
134
            export_paddle_param(param, name, py_code_dir)
J
jiangjiajun 已提交
135
        self.add_heads()
J
jiangjiajun 已提交
136 137 138 139

        if hasattr(self, "used_custom_layers"):
            for _, layer_code in self.used_custom_layers.items():
                self.add_codes(layer_code, 0)
J
jiangjiajun 已提交
140
                self.add_codes("", 0)
J
jiangjiajun 已提交
141 142 143 144 145

        self.add_codes("\ndef x2paddle_net():", 0)
        for i in range(len(self.graph.topo_sort)):
            node_name = self.graph.topo_sort[i]
            node = self.graph.get_node(node_name)
J
jiangjiajun 已提交
146 147
            if len(node.fluid_code.layers) == 0:
                continue
J
jiangjiajun 已提交
148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167
            self.add_codes(node.fluid_code.gen_codes(), 1)

        self.add_codes("", 0)

        input_str = "["
        for name in self.graph.input_nodes:
            input_str += (name + ", ")
        input_str = input_str.strip(", ") + "]"
        output_str = "["
        for name in self.graph.output_nodes:
            output_str += (name + ", ")
        output_str = output_str.strip(", ") + "]"

        return_code = "return {}, {}".format(input_str, output_str)

        self.add_codes(return_code, 1)
        self.add_codes("", 0)

        self.add_codes(inspect.getsourcelines(run_net)[0])
        fp = open(os.path.join(py_code_dir, "model.py"), 'w')
J
jiangjiajun 已提交
168 169
        fp.write(self.paddle_codes)
        fp.close()