op_mapper.py 7.4 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 58 59 60 61 62 63 64 65 66 67 68 69
# This func will copy to generate code file
def run_net(param_dir="./"):
    import os
    inputs, outputs = x2paddle_net()
    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

70 71
    fluid.io.load_vars(
        exe, param_dir, fluid.default_main_program(), predicate=if_exist)
C
channingss 已提交
72 73


J
jiangjiajun 已提交
74 75 76 77 78 79
class OpMapper(object):
    def __init__(self):
        self.paddle_codes = ""
        self.tab = "    "
        self.net_code = list()
        self.weights = dict()
J
jiangjiajun 已提交
80 81
        self.inputs = list()
        self.outputs = list()
J
jiangjiajun 已提交
82

J
jiangjiajun 已提交
83 84 85 86 87 88 89 90 91 92
    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 已提交
93 94
            print("There are {} ops not supported yet, list as below".format(
                len(unsupported_ops)))
J
jiangjiajun 已提交
95 96 97 98
            for op in unsupported_ops:
                print(op)
            return False

J
jiangjiajun 已提交
99 100 101
    def add_codes(self, codes, indent=0):
        if isinstance(codes, list):
            for code in codes:
102 103
                self.paddle_codes += (
                    self.tab * indent + code.strip('\n') + '\n')
J
jiangjiajun 已提交
104
        elif isinstance(codes, str):
J
jiangjiajun 已提交
105
            self.paddle_codes += (self.tab * indent + codes.strip('\n') + '\n')
J
jiangjiajun 已提交
106 107 108 109 110 111 112 113 114
        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 已提交
115
    def save_inference_model(self, save_dir, params_merge):
J
jiangjiajun 已提交
116 117 118 119 120 121 122 123 124
        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()
S
SunAhong1993 已提交
125 126 127 128 129
            for i, out in enumerate(outputs):
                if isinstance(out, list):
                    for out_part in out:
                        outputs.append(out_part)
                    del outputs[i]
J
jiangjiajun 已提交
130 131 132 133 134 135
            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 已提交
136
                    os.path.join(os.path.join(py_code_dir, var.name)))
J
jiangjiajun 已提交
137 138
                return b

139 140 141 142 143
            fluid.io.load_vars(
                exe,
                py_code_dir,
                fluid.default_main_program(),
                predicate=if_exist)
M
mamingjie-China 已提交
144
            if params_merge:
145 146 147 148 149 150
                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 已提交
151
            else:
152 153 154 155 156 157
                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 已提交
158 159 160 161
        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 已提交
162 163

    def save_python_model(self, save_dir):
J
jiangjiajun 已提交
164 165 166 167 168 169 170
        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 已提交
171
        for name, param in self.weights.items():
J
jiangjiajun 已提交
172
            export_paddle_param(param, name, py_code_dir)
J
jiangjiajun 已提交
173
        self.add_heads()
J
jiangjiajun 已提交
174 175 176 177

        if hasattr(self, "used_custom_layers"):
            for _, layer_code in self.used_custom_layers.items():
                self.add_codes(layer_code, 0)
J
jiangjiajun 已提交
178
                self.add_codes("", 0)
J
jiangjiajun 已提交
179 180 181 182 183

        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 已提交
184 185
            if node is None:
                continue
J
jiangjiajun 已提交
186 187
            if len(node.fluid_code.layers) == 0:
                continue
J
jiangjiajun 已提交
188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207
            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 已提交
208 209
        fp.write(self.paddle_codes)
        fp.close()