op_mapper.py 8.5 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.
C
channingss 已提交
14
import paddle.fluid as fluid
J
modify  
jiangjiajun 已提交
15
from paddle.fluid.proto import framework_pb2
J
jiangjiajun 已提交
16
from x2paddle.core.util import *
J
jiangjiajun 已提交
17
import inspect
J
jiangjiajun 已提交
18 19 20
import os


J
jiangjiajun 已提交
21 22 23 24 25 26 27
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'],
J
jiangjiajun 已提交
28 29
        "float64": [framework_pb2.VarType.FP64, 'd'],
        "bool": [framework_pb2.VarType.BOOL, None]
J
jiangjiajun 已提交
30 31
    }
    shape = param.shape
32 33
    if str(param.dtype) in ['uint8', 'uint_8', 'bool']:
        param = param.astype('int64')
J
jiangjiajun 已提交
34 35 36
    if len(shape) == 0:
        assert param.size == 1, "Unexpected situation happend!"
        shape = [1]
37 38 39
    assert str(
        param.dtype) in dtype_map, "Unknown dtype {} of params: {}.".format(
            str(param.dtype), param_name)
J
jiangjiajun 已提交
40 41 42 43 44 45 46 47 48 49 50 51 52 53
    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()


C
channingss 已提交
54 55 56 57
# This func will copy to generate code file
def run_net(param_dir="./"):
    import os
    inputs, outputs = x2paddle_net()
M
mamingjie-China 已提交
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75

    ops = fluid.default_main_program().global_block().ops
    used_vars = list()
    for op in ops:
        used_vars += op.input_arg_names

    tmp = list()
    for input in inputs:
        if isinstance(input, list):
            for ipt in input:
                if ipt.name not in used_vars:
                    continue
                tmp.append(ipt)
        else:
            if input.name not in used_vars:
                continue
            tmp.append(input)
    inputs = tmp
C
channingss 已提交
76 77 78 79 80 81 82 83 84 85 86 87
    for i, out in enumerate(outputs):
        if isinstance(out, list):
            for out_part in out:
                outputs.append(out_part)
            del outputs[i]
    exe = fluid.Executor(fluid.CPUPlace())
    exe.run(fluid.default_startup_program())

    def if_exist(var):
        b = os.path.exists(os.path.join(param_dir, var.name))
        return b

88 89
    fluid.io.load_vars(
        exe, param_dir, fluid.default_main_program(), predicate=if_exist)
C
channingss 已提交
90 91


J
jiangjiajun 已提交
92 93 94 95 96 97
class OpMapper(object):
    def __init__(self):
        self.paddle_codes = ""
        self.tab = "    "
        self.net_code = list()
        self.weights = dict()
J
jiangjiajun 已提交
98 99
        self.inputs = list()
        self.outputs = list()
J
jiangjiajun 已提交
100

J
jiangjiajun 已提交
101 102 103 104 105 106 107 108 109 110
    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 已提交
111 112
            print("There are {} ops not supported yet, list as below".format(
                len(unsupported_ops)))
J
jiangjiajun 已提交
113 114 115 116
            for op in unsupported_ops:
                print(op)
            return False

J
jiangjiajun 已提交
117 118 119
    def add_codes(self, codes, indent=0):
        if isinstance(codes, list):
            for code in codes:
120 121
                self.paddle_codes += (
                    self.tab * indent + code.strip('\n') + '\n')
J
jiangjiajun 已提交
122
        elif isinstance(codes, str):
J
jiangjiajun 已提交
123
            self.paddle_codes += (self.tab * indent + codes.strip('\n') + '\n')
J
jiangjiajun 已提交
124 125 126 127 128 129 130 131 132
        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("")

M
mamingjie-China 已提交
133
    def save_inference_model(self, save_dir, params_merge):
J
jiangjiajun 已提交
134 135 136 137 138 139 140 141 142
        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()
M
mamingjie-China 已提交
143 144 145 146 147 148

            ops = fluid.default_main_program().global_block().ops
            used_vars = list()
            for op in ops:
                used_vars += op.input_arg_names

S
SunAhong1993 已提交
149 150 151 152 153
            for i, out in enumerate(outputs):
                if isinstance(out, list):
                    for out_part in out:
                        outputs.append(out_part)
                    del outputs[i]
M
mamingjie-China 已提交
154 155 156 157 158 159 160 161 162 163 164 165 166

            input_names = list()
            for input in inputs:
                if isinstance(input, list):
                    for ipt in input:
                        if ipt.name not in used_vars:
                            continue
                        input_names.append(ipt.name)
                else:
                    if input.name not in used_vars:
                        continue
                    input_names.append(input.name)

J
jiangjiajun 已提交
167 168 169 170 171
            exe = fluid.Executor(fluid.CPUPlace())
            exe.run(fluid.default_startup_program())

            def if_exist(var):
                b = os.path.exists(
J
jiangjiajun 已提交
172
                    os.path.join(os.path.join(py_code_dir, var.name)))
J
jiangjiajun 已提交
173 174
                return b

175 176 177 178 179
            fluid.io.load_vars(
                exe,
                py_code_dir,
                fluid.default_main_program(),
                predicate=if_exist)
M
mamingjie-China 已提交
180
            if params_merge:
181 182 183 184 185 186
                fluid.io.save_inference_model(
                    dirname=os.path.join(save_dir, "inference_model"),
                    feeded_var_names=input_names,
                    target_vars=outputs,
                    executor=exe,
                    params_filename="__params__")
M
mamingjie-China 已提交
187
            else:
188 189 190 191 192 193
                fluid.io.save_inference_model(
                    dirname=os.path.join(save_dir, "inference_model"),
                    feeded_var_names=input_names,
                    target_vars=outputs,
                    executor=exe,
                    params_filename=None)
J
jiangjiajun 已提交
194 195 196 197
        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 已提交
198 199

    def save_python_model(self, save_dir):
J
jiangjiajun 已提交
200 201 202 203 204 205 206
        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 已提交
207
        for name, param in self.weights.items():
J
jiangjiajun 已提交
208
            export_paddle_param(param, name, py_code_dir)
J
jiangjiajun 已提交
209
        self.add_heads()
J
jiangjiajun 已提交
210 211 212 213

        if hasattr(self, "used_custom_layers"):
            for _, layer_code in self.used_custom_layers.items():
                self.add_codes(layer_code, 0)
J
jiangjiajun 已提交
214
                self.add_codes("", 0)
J
jiangjiajun 已提交
215 216 217 218 219

        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 已提交
220 221
            if node is None:
                continue
J
jiangjiajun 已提交
222 223
            if len(node.fluid_code.layers) == 0:
                continue
J
jiangjiajun 已提交
224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243
            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 已提交
244 245
        fp.write(self.paddle_codes)
        fp.close()