tf_op_mapper_nhwc.py 44.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
#   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.

from x2paddle.decoder.tf_decoder import TFGraph
from x2paddle.core.op_mapper import OpMapper
from x2paddle.core.util import *
J
jiangjiajun 已提交
18 19 20
from x2paddle import program
from x2paddle import gen_name
import traceback
M
mamingjie-China 已提交
21
import math
22 23 24 25 26 27 28 29 30
import inspect
import numpy
import sys


# compute padding size for SAME mode
def get_same_padding(in_size, kernel_size, stride):
    new_size = int(math.ceil(in_size * 1.0 / stride))
    pad_size = (new_size - 1) * stride + kernel_size - in_size
J
jiangjiajun 已提交
31 32
    if pad_size < 0:
        pad_size = 0
33 34 35 36 37 38 39 40 41 42 43 44 45
    pad0 = int(pad_size / 2)
    pad1 = pad_size - pad0
    return [pad0, pad1]


class TFOpMapperNHWC(OpMapper):
    directly_map_ops = {
        'Relu': ['relu'],
        'Relu6': ['relu6'],
        'Abs': ['abs'],
        'Sigmoid': ['sigmoid'],
        'Exp': ['exp'],
        'Rsqrt': ['rsqrt'],
J
jiangjiajun@baidu.com 已提交
46
        'Sqrt': ['sqrt'],
47
        'swish_f32': ['swish'],
48
        'Tanh': ['tanh'],
J
jiangjiajun 已提交
49
        'Softplus': ['softplus'],
50 51
        'LeakyRelu': ['leaky_relu', {
            'alpha': 'alpha'
52
        }],
M
mamingjie-China 已提交
53
        'Floor': ['floor'],
J
jiangjiajun 已提交
54 55
        'Erf': ['erf'],
        'Square': ['square']
56 57 58
    }
    elementwise_ops = {
        'Add': 'elementwise_add',
J
jiangjiajun@baidu.com 已提交
59
        'AddV2': 'elementwise_add',
60 61 62
        'RealDiv': 'elementwise_div',
        'Sub': 'elementwise_sub',
        'Maximum': 'elementwise_max',
63
        'Minimum': 'elementwise_min',
M
mamingjie-China 已提交
64
        'LessEqual': 'less_equal',
J
jiangjiajun 已提交
65
        'GreaterEqual': 'greater_equal',
J
jiangjiajun 已提交
66 67
        'Mul': 'elementwise_mul',
        'FloorDiv': 'elementwise_floordiv'
68 69 70 71 72 73 74 75 76
    }

    def __init__(self, decoder):
        super(TFOpMapperNHWC, self).__init__()
        self.decoder = decoder
        self.graph = decoder.tf_graph
        self.weights = dict()
        self.omit_nodes = list()
        self.used_custom_layers = dict()
J
jiangjiajun 已提交
77
        program.clear()
78 79 80

        not_placeholder = list()
        for name in self.graph.input_nodes:
M
mamingjie-China 已提交
81 82 83 84 85
            if self.graph.get_node(
                    name).layer_type != "Placeholder" and self.graph.get_node(
                        name
                    ).layer_type != "OneShotIterator" and self.graph.get_node(
                        name).layer_type != "IteratorV2":
86 87 88 89 90
                not_placeholder.append(name)
        for name in not_placeholder:
            idx = self.graph.input_nodes.index(name)
            del self.graph.input_nodes[idx]

J
jiangjiajun 已提交
91 92 93
        program.inputs = self.graph.input_nodes
        program.outputs = self.graph.output_nodes

94
        unsupported_ops = set()
95 96
        sys.stderr.write("Total nodes: {}\n".format(len(self.graph.topo_sort)))
        for i, node_name in enumerate(self.graph.topo_sort):
M
mamingjie-China 已提交
97
            sys.stderr.write("\rConverting node {} ...     ".format(i + 1))
98 99 100 101 102 103 104 105 106 107 108 109 110 111
            node = self.graph.get_node(node_name)
            op = node.layer_type
            if op in self.directly_map_ops:
                if len(unsupported_ops) > 0:
                    continue
                self.directly_map(node)
            elif op in self.elementwise_ops:
                if len(unsupported_ops) > 0:
                    continue
                self.elementwise_map(node)
            elif hasattr(self, op):
                if len(unsupported_ops) > 0:
                    continue
                func = getattr(self, op)
J
jiangjiajun@baidu.com 已提交
112 113
                try:
                    func(node)
114
                except Exception as e:
J
jiangjiajun@baidu.com 已提交
115
                    unsupported_ops.add(op)
J
jiangjiajun 已提交
116
                    print("\n{}\n".format(traceback.format_exc()))
117 118 119
            else:
                unsupported_ops.add(op)
        if len(unsupported_ops) > 0:
J
jiangjiajun 已提交
120
            print("\n========= {} OPs are not supported yet ===========".format(
121 122 123 124
                len(unsupported_ops)))
            for op in unsupported_ops:
                print("========== {} ============".format(op))
            sys.exit(-1)
M
mamingjie-China 已提交
125
        sys.stderr.write("\nDone!\n")
126 127 128 129

    def directly_map(self, node):
        assert node.layer_type in self.directly_map_ops
        op_info = self.directly_map_ops[node.layer_type]
J
jiangjiajun 已提交
130
        input = self.graph.get_node(node.layer.input[0])
131 132 133 134 135 136
        attr = dict()
        for param in op_info[1:]:
            tf_param_name = list(param.keys())[0]
            pd_param_name = list(param.values())[0]
            tf_param = node.get_attr(tf_param_name)
            attr[pd_param_name] = tf_param
M
modify  
mamingjie-China 已提交
137

J
jiangjiajun 已提交
138 139 140 141 142
        program.add_layer(
            kernel="fluid.layers.{}".format(op_info[0]),
            inputs={"x": input.name},
            outputs=[node.name],
            **attr)
143 144 145 146

    def elementwise_map(self, node):
        assert node.layer_type in self.elementwise_ops
        op_type = self.elementwise_ops[node.layer_type]
J
jiangjiajun 已提交
147 148
        x = self.graph.get_node(node.layer.input[0])
        y = self.graph.get_node(node.layer.input[1])
J
jiangjiajun 已提交
149

J
jiangjiajun 已提交
150 151 152 153 154
        program.add_layer(
            kernel="fluid.layers.{}".format(op_type),
            inputs={"x": x.name,
                    "y": y.name},
            outputs=[node.name])
155

J
jiangjiajun 已提交
156 157 158 159 160 161 162 163 164 165
    def NotEqual(self, node):
        x = self.graph.get_node(node.layer.input[0])
        y = self.graph.get_node(node.layer.input[1])

        program.add_layer(
            kernel="fluid.layers.not_equal",
            inputs={"x": x.name,
                    "y": y.name},
            outputs=[node.name])

166 167 168 169 170
    def Placeholder(self, node):
        shape = node.out_shapes[0]
        assert len(shape) != 0, "Unknown shape of input nodes[{}].".format(
            node.layer_name)
        dtype = node.dtype
J
jiangjiajun 已提交
171 172 173 174 175 176 177
        program.add_layer(
            kernel="fluid.data",
            inputs={},
            outputs=[node.name],
            dtype=string(dtype),
            shape=shape,
            name=string(node.name))
178 179 180 181 182 183 184 185 186

    def Const(self, node):
        shape = node.out_shapes[0]
        dtype = node.dtype
        value = node.value
        initializer = "Constant(0.0)"
        if len(shape) == 0:
            assert value.size == 1, "Unexpected situation happend"
            shape = [1]
J
jiangjiajun 已提交
187 188
            if value == float('inf'):
                value = "float('inf')"
189 190
            initializer = "Constant({})".format(value)

J
jiangjiajun 已提交
191 192 193 194 195 196 197 198 199
        program.parameters[node.name] = node.value
        program.add_layer(
            kernel="fluid.layers.create_parameter",
            inputs={},
            outputs=[node.name],
            dtype=string(dtype),
            shape=shape,
            name=string(node.name),
            default_initializer=initializer)
200 201

    def Transpose(self, node):
J
jiangjiajun 已提交
202 203
        input = self.graph.get_node(node.layer.input[0])
        perm = self.graph.get_node(node.layer.input[1])
204 205 206
        assert perm.layer_type == "Const", "Perm of transpose OP should be Const"
        perm = perm.value.tolist()

J
jiangjiajun 已提交
207 208 209 210 211
        program.add_layer(
            kernel="fluid.layers.transpose",
            inputs={"x": input.name},
            outputs=[node.name],
            perm=perm)
212

213
    def Fill(self, node):
M
update  
mamingjie-China 已提交
214 215
        dims = self.graph.get_node(node.layer.input[0])
        input_value = self.graph.get_node(node.layer.input[1])
M
update  
mamingjie-China 已提交
216 217
        inputs = dict()
        attr = dict()
218
        assert input_value.layer_type == "Const", "Value of fill OP should be Const"
M
update  
mamingjie-China 已提交
219 220 221 222 223 224
        if dims.layer_type == "Const":
            attr["shape"] = dims.value.tolist()
        else:
            inputs["shape"] = dims.name
        attr["dtype"] = string(input_value.dtype)
        attr["value"] = input_value.value
225

J
jiangjiajun 已提交
226 227
        program.add_layer(
            "fluid.layers.fill_constant",
M
update  
mamingjie-China 已提交
228
            inputs=inputs,
J
jiangjiajun 已提交
229
            outputs=[node.name],
M
update  
mamingjie-China 已提交
230
            **attr)
231 232

    def DepthToSpace(self, node):
M
update  
mamingjie-China 已提交
233
        input = self.graph.get_node(node.layer.input[0])
234 235 236

        block_size = node.get_attr("block_size")
        data_format = node.get_attr("data_format").decode()
M
update  
mamingjie-China 已提交
237 238 239 240
        if data_format == "NHWC":
            n, h, w, c = input.out_shapes[0]
        else:
            n, c, h, w = input.out_shapes[0]
241

J
jiangjiajun 已提交
242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274
        input_name = input.name
        if data_format == "NHWC":
            transpose_name = gen_name("depth_to_space", "transpose")
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": input.name},
                outputs=[transpose_name],
                perm=[0, 3, 1, 2])
            input_name = transpose_name

        shape = [0, block_size * block_size, -1, h, w]
        reshape_name = gen_name("depth_to_space", "reshape")
        program.add_layer(
            kernel="fluid.layers.reshape",
            inputs={"x": input_name},
            outputs=[reshape_name],
            shape=shape)

        transpose_name = gen_name("depth_to_space", "transpose")
        program.add_layer(
            kernel="fluid.layers.transpose",
            inputs={"x": reshape_name},
            outputs=[transpose_name],
            perm=[0, 2, 1, 3, 4])

        reshape_name = gen_name("depth_to_space", "reshape")
        program.add_layer(
            kernel="fluid.layers.reshape",
            inputs={"x": transpose_name},
            outputs=[reshape_name],
            shape=[0, c, h, w])

        program.add_layer(
M
update  
mamingjie-China 已提交
275 276
            kernel="fluid.layers.pixel_shuffle",
            inputs={"x": reshape_name},
J
jiangjiajun 已提交
277 278
            outputs=[node.name],
            upscale_factor=block_size)
279 280

        if data_format == "NHWC":
J
jiangjiajun 已提交
281 282 283 284 285
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": node.name},
                outputs=[node.name],
                perm=[0, 2, 3, 1])
286

287
    def MaxPool(self, node):
M
update  
mamingjie-China 已提交
288
        input = self.graph.get_node(node.layer.input[0])
289 290 291 292 293 294

        k_size = node.get_attr("ksize")
        strides = node.get_attr("strides")
        data_format = node.get_attr("data_format").decode()
        pad_mode = node.get_attr("padding").decode()

J
jiangjiajun 已提交
295 296 297 298 299 300 301 302
        input_name = input.name
        if data_format == "NHWC":
            transpose_name = gen_name("max_pool", "transpose")
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": input.name},
                outputs=[transpose_name],
                perm=[0, 3, 1, 2])
303 304
            strides = [strides[i] for i in [0, 3, 1, 2]]
            k_size = [k_size[i] for i in [0, 3, 1, 2]]
J
jiangjiajun 已提交
305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321
            input_name = transpose_name

        program.add_layer(
            kernel="fluid.layers.pool2d",
            inputs={"input": input_name},
            outputs=[node.name],
            pool_size=k_size[2:4],
            pool_type=string("max"),
            pool_stride=strides[2:4],
            pool_padding=string(pad_mode))

        if data_format == "NHWC":
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": node.name},
                outputs=[node.name],
                perm=[0, 2, 3, 1])
322 323

    def Conv2D(self, node):
J
jiangjiajun 已提交
324 325
        input = self.graph.get_node(node.layer.input[0])
        kernel = self.graph.get_node(node.layer.input[1])
326 327 328 329 330 331

        k_size = kernel.out_shapes[0]
        strides = node.get_attr("strides")
        dilations = node.get_attr("dilations")
        data_format = node.get_attr("data_format").decode()
        pad_mode = node.get_attr("padding").decode()
M
mamingjie-China 已提交
332 333
        if data_format == "NHWC":
            n, h, w, c = input.out_shapes[0]
M
update  
mamingjie-China 已提交
334 335
        else:
            n, c, h, w = input.out_shapes[0]
336

J
jiangjiajun@baidu.com 已提交
337 338
        if kernel.layer_type == 'Const':
            kernel_value = kernel.value
M
mamingjie-China 已提交
339
            kernel_weight_name = kernel.name.replace('/', '_')
340 341 342
        else:
            kernel_value = self.decoder.infer_tensor(kernel)
            if kernel.layer_type == 'Split':
M
mamingjie-China 已提交
343 344
                kernel_weight_name = "{}_{}_kernel".format(node.name,
                                                           kernel.name)
345
            else:
M
mamingjie-China 已提交
346
                kernel_weight_name = kernel.name.replace('/', '_')
J
jiangjiajun 已提交
347 348
        program.parameters[kernel_weight_name] = numpy.transpose(kernel_value,
                                                                 (3, 2, 0, 1))
349

J
jiangjiajun 已提交
350 351
        input_name = input.name
        if data_format == "NHWC":
352 353
            strides = [strides[i] for i in [0, 3, 1, 2]]
            dilations = [dilations[i] for i in [0, 3, 1, 2]]
J
jiangjiajun 已提交
354 355 356 357 358 359 360 361
            transpose_name = gen_name("conv2d", "transpose")
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": input.name},
                outputs=[transpose_name],
                perm=[0, 3, 1, 2])
            input_name = transpose_name

M
mamingjie-China 已提交
362 363 364 365 366 367 368 369 370 371
        if c == -1:
            attr = {"shape": [0, k_size[2], 0, 0]}
            node.fluid_code.add_layer(
                "reshape", inputs=input, output=input, param_attr=attr)
            program.add_layer(
                kernel="fluid.layers.reshape",
                inputs={"x": input_name},
                outputs=[input_name],
                shape=[0, k_size[2], 0, 0])

J
jiangjiajun 已提交
372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389
        program.add_layer(
            kernel="fluid.layers.conv2d",
            inputs={"input": input_name},
            outputs=[node.name],
            bias_attr=False,
            param_attr=string(kernel_weight_name),
            num_filters=k_size[3],
            filter_size=k_size[0:2],
            stride=strides[2:4],
            dilation=dilations[2:4],
            padding=string(pad_mode))

        if data_format == "NHWC":
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": node.name},
                outputs=[node.name],
                perm=[0, 2, 3, 1])
390 391

    def BiasAdd(self, node):
J
jiangjiajun 已提交
392 393 394 395 396 397 398
        input = self.graph.get_node(node.layer.input[0])
        bias = self.graph.get_node(node.layer.input[1])
        program.add_layer(
            kernel="fluid.layers.elementwise_add",
            inputs={"x": input.name,
                    "y": bias.name},
            outputs=[node.name])
399 400

    def FusedBatchNorm(self, node):
J
jiangjiajun 已提交
401 402 403 404 405
        input = self.graph.get_node(node.layer.input[0])
        gamma = self.graph.get_node(node.layer.input[1])
        beta = self.graph.get_node(node.layer.input[2])
        moving_mean = self.graph.get_node(node.layer.input[3])
        moving_var = self.graph.get_node(node.layer.input[4])
406 407 408 409 410 411 412
        data_format = node.get_attr("data_format").decode()

        assert gamma.layer_type == "Const"
        assert beta.layer_type == "Const"
        assert moving_mean.layer_type == "Const"
        assert moving_var.layer_type == "Const"

J
jiangjiajun 已提交
413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432
        input_name = input.name
        if data_format == "NHWC":
            transpose_name = gen_name("batch_norm", "transpose")
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": input.name},
                outputs=[transpose_name],
                perm=[0, 3, 1, 2])
            input_name = transpose_name

        program.add_layer(
            kernel="fluid.layers.batch_norm",
            inputs={"input": input_name},
            outputs=[node.name],
            epsilon=node.get_attr("epsilon"),
            param_attr=string(gamma.name),
            bias_attr=string(beta.name),
            moving_mean_name=string(moving_mean.name),
            moving_variance_name=string(moving_var.name),
            is_test=True)
433

J
jiangjiajun 已提交
434 435 436 437 438 439
        if data_format == "NHWC":
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": node.name},
                outputs=[node.name],
                perm=[0, 2, 3, 1])
440

J
jiangjiajun 已提交
441 442 443 444 445 446
    def Mean(self, node):
        input = self.graph.get_node(node.layer.input[0])
        reduce_idx = self.graph.get_node(node.layer.input[1])
        assert reduce_idx.layer_type == "Const", "Only support Const parameter[reduce_idx]"
        dims = reduce_idx.value.tolist()
        keep_dims = node.get_attr("keep_dims")
447

J
jiangjiajun 已提交
448 449 450 451 452 453
        program.add_layer(
            kernel="fluid.layers.reduce_mean",
            inputs={"input": input.name},
            outputs=[node.name],
            dim=dims,
            keep_dim=keep_dims)
454 455

    def Reshape(self, node):
J
jiangjiajun 已提交
456 457
        input = self.graph.get_node(node.layer.input[0])
        param = self.graph.get_node(node.layer.input[1])
J
jiangjiajun 已提交
458 459 460 461 462 463 464 465 466 467 468

        input_name = input.name
        if input.dtype == 'bool':
            cast_name = gen_name('reshape', 'cast')
            program.add_layer(
                kernel="fluid.layers.cast",
                inputs={"x": input_name},
                outputs=[cast_name],
                dtype="'int32'")
            input_name = cast_name

469
        if param.layer_type == "Const":
470
            shape = param.value.tolist()
J
jiangjiajun 已提交
471 472
            program.add_layer(
                kernel="fluid.layers.reshape",
J
jiangjiajun 已提交
473
                inputs={"x": input_name},
J
jiangjiajun 已提交
474 475
                outputs=[node.name],
                shape=shape)
476
        else:
J
jiangjiajun 已提交
477 478
            program.add_layer(
                kernel="fluid.layers.reshape",
J
jiangjiajun 已提交
479
                inputs={"x": input_name,
J
jiangjiajun 已提交
480 481
                        "shape": param.name},
                outputs=[node.name])
482 483 484 485
        if param.layer_type != "Const":
            out_shape = numpy.array(node.out_shapes[0])
            if (out_shape > 0).any():
                out_shape[out_shape < 0] = 0
J
jiangjiajun 已提交
486 487 488 489 490
                program.add_layer(
                    kernel="fluid.layers.reshape",
                    inputs={"x": node.name},
                    outputs=[node.name],
                    shape=out_shape.tolist())
491

J
jiangjiajun 已提交
492 493 494 495 496 497 498
        if input.dtype == 'bool':
            program.add_layer(
                kernel="fluid.layers.cast",
                inputs={"x": node.name},
                outputs=[node.name],
                dtype="'bool'")

499
    def Pad(self, node):
J
jiangjiajun 已提交
500 501
        input = self.graph.get_node(node.layer.input[0])
        paddings = self.graph.get_node(node.layer.input[1])
502 503 504 505
        assert paddings.layer_type == "Const", "Padding should be Const"
        paddings = paddings.value.flatten().tolist()

        if len(input.out_shapes[0]) == 4:
J
jiangjiajun 已提交
506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523
            if paddings[0] + paddings[1] + paddings[6] + paddings[7] == 0:
                new_padding = paddings[2:6]
                transpose_name = gen_name("pad", "transpose")
                program.add_layer(
                    kernel="fluid.layers.transpose",
                    inputs={"x": input.name},
                    outputs=[transpose_name],
                    perm=[0, 3, 1, 2])
                program.add_layer(
                    kernel="fluid.layers.pad2d",
                    inputs={"input": transpose_name},
                    outputs=[node.name],
                    paddings=new_padding)
                program.add_layer(
                    kernel="fluid.layers.transpose",
                    inputs={"x": node.name},
                    outputs=[node.name],
                    perm=[0, 2, 3, 1])
524 525
                return

J
jiangjiajun 已提交
526 527 528 529 530
        program.add_layer(
            kernel="fluid.layers.pad",
            inputs={"input": input.name},
            outputs=[node.name],
            paddings=paddings)
531

J
jiangjiajun 已提交
532 533 534 535 536 537 538 539
    def Squeeze(self, node):
        input = self.graph.get_node(node.layer.input[0])
        squeeze_dims = node.get_attr('squeeze_dims')
        program.add_layer(
            kernel="fluid.layers.squeeze",
            inputs={"input": input.name},
            outputs=[node.name],
            axes=squeeze_dims)
540

J
jiangjiajun 已提交
541 542 543 544 545 546 547 548 549 550 551
    def Softmax(self, node):
        input = self.graph.get_node(node.layer.input[0])
        axis = node.get_attr("axis")
        program.add_layer(
            kernel="fluid.layers.softmax",
            inputs={"input": input.name},
            outputs=[node.name],
            axis=axis)

    def Shape(self, node):
        input = self.graph.get_node(node.layer.input[0])
J
jiangjiajun 已提交
552 553 554 555 556 557 558 559 560
        input_name = input.name
        if input.dtype == 'bool':
            cast_name = gen_name('shape', 'cast')
            program.add_layer(
                kernel="fluid.layers.cast",
                inputs={"x": input.name},
                outputs=[cast_name],
                dtype="'int32'")
            input_name = cast_name
J
jiangjiajun 已提交
561 562
        program.add_layer(
            kernel="fluid.layers.shape",
J
jiangjiajun 已提交
563
            inputs={"input": input_name},
J
jiangjiajun 已提交
564
            outputs=[node.name])
565

J
jiangjiajun 已提交
566 567 568 569 570 571 572 573 574 575
    def ArgMax(self, node):
        input = self.graph.get_node(node.layer.input[0])
        axis = self.graph.get_node(node.layer.input[1])
        assert axis.layer_type == "Const", "ArgMax only support Const parameter"
        axis = axis.value
        program.add_layer(
            kernel="fluid.layers.argmax",
            inputs={"x": input.name},
            outputs=[node.name],
            axis=axis)
576 577

    def MatMul(self, node):
J
jiangjiajun 已提交
578 579
        x = self.graph.get_node(node.layer.input[0])
        y = self.graph.get_node(node.layer.input[1])
580 581
        transpose_a = node.get_attr('transpose_a')
        transpose_b = node.get_attr('transpose_b')
M
mamingjie-China 已提交
582 583 584 585
        if transpose_a is None:
            transpose_a = node.get_attr('adj_x')
        if transpose_b is None:
            transpose_b = node.get_attr('adj_y')
J
jiangjiajun 已提交
586 587 588 589 590 591 592 593
        program.add_layer(
            kernel="fluid.layers.matmul",
            inputs={"x": x.name,
                    "y": y.name},
            outputs=[node.name],
            transpose_x=transpose_a,
            transpose_y=transpose_b)

M
mamingjie-China 已提交
594 595 596 597 598 599
    def BatchMatMul(self, node):
        return self.MatMul(node)

    def BatchMatMulV2(self, node):
        return self.MatMul(node)

J
jiangjiajun 已提交
600 601 602 603 604 605 606 607 608 609 610
    def DepthwiseConv2dNative(self, node):
        input = self.graph.get_node(node.layer.input[0])
        kernel = self.graph.get_node(node.layer.input[1])
        assert kernel.layer_type == "Const", "Kernel of DepthwiseConv2DNative should be Const"

        in_shape = input.out_shapes[0]
        k_size = kernel.out_shapes[0]
        strides = node.get_attr("strides")
        dilations = node.get_attr("dilations")
        data_format = node.get_attr("data_format").decode()
        pad_mode = node.get_attr("padding").decode()
611

J
jiangjiajun 已提交
612 613
        program.parameters[kernel.layer_name.replace(
            '/', '_')] = numpy.transpose(kernel.value, (2, 3, 0, 1))
M
mamingjie-China 已提交
614

J
jiangjiajun 已提交
615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639
        input_name = input.name
        if data_format == "NHWC":
            in_shape = [in_shape[i] for i in [0, 3, 1, 2]]
            strides = [strides[i] for i in [0, 3, 1, 2]]
            dilations = [dilations[i] for i in [0, 3, 1, 2]]
            transpose_name = gen_name('depthwise_conv2d', 'transpose')
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": input.name},
                outputs=[transpose_name],
                perm=[0, 3, 1, 2])
            input_name = transpose_name

        program.add_layer(
            kernel="fluid.layers.conv2d",
            inputs={"input": input_name},
            outputs=[node.name],
            num_filters=in_shape[1],
            filter_size=k_size[0:2],
            stride=strides[2:4],
            dilation=dilations[2:4],
            groups=k_size[3] * in_shape[1],
            padding=string(pad_mode),
            param_attr=string(kernel.layer_name),
            bias_attr=False)
M
mamingjie-China 已提交
640

J
jiangjiajun 已提交
641 642 643 644 645 646 647 648
        if data_format == "NHWC":
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": node.name},
                outputs=[node.name],
                perm=[0, 2, 3, 1])

    def AvgPool(self, node):
M
update  
mamingjie-China 已提交
649
        input = self.graph.get_node(node.layer.input[0])
J
jiangjiajun 已提交
650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685

        k_size = node.get_attr("ksize")
        strides = node.get_attr("strides")
        data_format = node.get_attr("data_format").decode()
        pad_mode = node.get_attr("padding").decode()

        input_name = input.name
        if data_format == "NHWC":
            transpose_name = gen_name("avg_pool", "transpose")
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": input.name},
                outputs=[transpose_name],
                perm=[0, 3, 1, 2])
            strides = [strides[i] for i in [0, 3, 1, 2]]
            k_size = [k_size[i] for i in [0, 3, 1, 2]]
            input_name = transpose_name

        program.add_layer(
            kernel="fluid.layers.pool2d",
            inputs={"input": input_name},
            outputs=[node.name],
            pool_size=k_size[2:4],
            pool_type=string("avg"),
            pool_stride=strides[2:4],
            pool_padding=string(pad_mode))

        if data_format == "NHWC":
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": node.name},
                outputs=[node.name],
                perm=[0, 2, 3, 1])

    def Pack(self, node):
        inputs = [self.graph.get_node(name) for name in node.layer.input]
J
jiangjiajun 已提交
686
        input_names = [i.name for i in inputs]
J
jiangjiajun 已提交
687 688 689
        axis = node.get_attr("axis")
        program.add_layer(
            kernel="fluid.layers.stack",
J
jiangjiajun 已提交
690
            inputs={"x": input_names},
J
jiangjiajun 已提交
691 692
            outputs=[node.name],
            axis=axis)
J
jiangjiajun 已提交
693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722
        if len(node.out_shapes[0]) == 1:
            program.add_layer(
                kernel="fluid.layers.reshape",
                inputs={"x": node.name},
                outputs=[node.name],
                shape=[-1])

    def Unpack(self, node):
        input = self.graph.get_node(node.layer.input[0])
        axis = node.get_attr("axis")
        num = node.get_attr("num")
        shape = input.out_shapes[0]
        input_name = input.name
        if len(shape) == 1:
            if shape[0] > 0 and num == shape[0]:
                program.add_layer(
                    kernel="fluid.layers.unsqueeze",
                    inputs={"input": input.name},
                    outputs=[node.name],
                    axes=[0])
                input_name = node.name
                axis = 1
            else:
                raise Exception("Unexpected situation happend in Unpack OP")
        program.add_layer(
            kernel="fluid.layers.unstack",
            inputs={"x": input_name},
            outputs=["{}_p{}".format(node.layer_name, i) for i in range(num)],
            axis=axis,
            num=num)
J
jiangjiajun 已提交
723 724 725 726 727

    def ConcatV2(self, node):
        inputs = [self.graph.get_node(name) for name in node.layer.input[:-1]]
        axis = self.graph.get_node(node.layer.input[-1])
        assert axis.layer_type == "Const", "axis for ConcatV2 must be type Const"
728
        axis = axis.value
J
jiangjiajun 已提交
729 730
        if axis < 0:
            axis += len(inputs[0].out_shapes[0])
J
jiangjiajun 已提交
731 732 733 734 735 736 737 738 739 740 741

        input_names = [i.name for i in inputs]
        for i, ipt in enumerate(inputs):
            if node.dtype == 'bool':
                cast_name = gen_name('concat', 'cast')
                program.add_layer(
                    kernel="fluid.layers.cast",
                    inputs={"x": ipt.name},
                    outputs=[cast_name],
                    dtype="'int32'")
                input_names[i] = cast_name
J
jiangjiajun 已提交
742 743
        program.add_layer(
            kernel="fluid.layers.concat",
J
jiangjiajun 已提交
744
            inputs={"input": input_names},
J
jiangjiajun 已提交
745 746
            outputs=[node.name],
            axis=axis)
J
jiangjiajun 已提交
747 748 749 750 751 752
        if node.dtype == 'bool':
            program.add_layer(
                kernel="fluid.layers.cast",
                inputs={"x": node.name},
                outputs=[node.name],
                dtype="'bool'")
753 754

    def StridedSlice(self, node):
J
jiangjiajun 已提交
755 756 757 758
        input = self.graph.get_node(node.layer.input[0])
        begin = self.graph.get_node(node.layer.input[1])
        end = self.graph.get_node(node.layer.input[2])
        strides = self.graph.get_node(node.layer.input[3])
J
jiangjiajun 已提交
759 760 761 762 763 764 765 766 767 768 769 770 771 772

        if strides.layer_type == "Const":
            strides = strides.value.tolist()
        else:
            strides = self.decoder.infer_shape_tensor(strides)
        if begin.layer_type == "Const":
            begin = begin.value.tolist()
        else:
            begin = self.decoder.infer_shape_tensor(begin)
        if end.layer_type == "Const":
            end = end.value.tolist()
        else:
            end = self.decoder.infer_shape_tensor(end)

773 774 775
        assert len(set(strides)) == 1 and strides[
            0] == 1, "Only support strides be 1 in StridedSlice OP"

J
jiangjiajun 已提交
776 777 778 779
        if len(begin) < len(input.out_shapes[0]):
            begin = begin + [0] * (len(input.out_shapes[0]) - len(begin))
        if len(end) < len(input.out_shapes[0]):
            end = end + [0] * (len(input.out_shapes[0]) - len(end))
780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820
        for i in range(len(end)):
            if end[i] == 0:
                end[i] = 999999

        begin_mask = node.get_attr('begin_mask')
        end_mask = node.get_attr('end_mask')
        ellipsis_mask = node.get_attr('ellipsis_mask')
        new_axis_mask = node.get_attr('new_axis_mask')
        shrink_axis_mask = node.get_attr('shrink_axis_mask')

        assert ellipsis_mask == 0, "(OP:{} Name:{})Only support ellipsis_mask be 0[now: {}] n StridedSlice OP".format(
            node.layer_type, node.layer.name, ellipsis_mask)

        # TODO codes without validation
        # Use it carefully
        new_begin = list()
        new_end = list()
        new_axes = list()
        shrink_axes = list()
        for i, item in enumerate(begin):
            mask = (new_axis_mask >> i) & 1
            if mask != 0:
                new_axes.append(i)
                continue

            mask = (shrink_axis_mask >> i) & 1
            if mask != 0:
                shrink_axes.append(i)

            mask = (begin_mask >> i) & 1
            if mask != 0:
                new_begin.append(0)
            else:
                new_begin.append(item)

            mask = (end_mask >> i) & 1
            if mask != 0:
                new_end.append(999999)
            else:
                new_end.append(end[i])

J
jiangjiajun 已提交
821 822 823 824 825 826 827
        program.add_layer(
            kernel="fluid.layers.slice",
            inputs={"input": input.name},
            outputs=[node.name],
            axes=[i for i in range(len(new_begin))],
            starts=new_begin,
            ends=new_end)
828
        if len(new_axes) > 0:
J
jiangjiajun 已提交
829 830
            program.add_layer(
                kernel="fluid.layers.unsqueeze",
M
mamingjie-China 已提交
831
                inputs={"input": node.name},
J
jiangjiajun 已提交
832 833
                outputs=[node.name],
                axes=new_axes)
834 835 836 837
        if len(shrink_axes) > 0:
            if len(input.out_shapes[0]) + len(new_axes) <= 1:
                pass
            else:
J
jiangjiajun 已提交
838
                program.add_layer(
J
jiangjiajun 已提交
839
                    kernel="fluid.layers.squeeze",
M
mamingjie-China 已提交
840
                    inputs={"input": node.name},
J
jiangjiajun 已提交
841
                    outputs=[node.name],
J
jiangjiajun 已提交
842
                    axes=shrink_axes)
J
jiangjiajun 已提交
843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858

    def Split(self, node):
        dim = self.graph.get_node(node.layer.input[0])
        input = self.graph.get_node(node.layer.input[1])
        assert dim.layer_type == "Const"
        num_split = node.get_attr('num_split')
        dim = dim.value

        program.add_layer(
            kernel="fluid.layers.split",
            inputs={"input": input.name},
            outputs=[
                "{}_p{}".format(node.layer_name, i) for i in range(num_split)
            ],
            num_or_sections=num_split,
            dim=dim)
859 860

    def Slice(self, node):
J
jiangjiajun 已提交
861 862 863 864 865 866
        input = self.graph.get_node(node.layer.input[0])
        begin = self.graph.get_node(node.layer.input[1])
        size = self.graph.get_node(node.layer.input[2])

        inputs = {"x": input.name}
        attrs = {}
867 868
        if begin.layer_type == "Const":
            begin = begin.value.tolist()
J
jiangjiajun 已提交
869
            attrs['offsets'] = begin
870
        else:
M
mamingjie-China 已提交
871 872 873 874 875 876 877 878 879 880
            #             shape = begin.out_shapes[0]
            #             reshape_name = gen_name("slice", "reshape")
            #             program.add_layer(
            #                 kernel="fluid.layers.reshape",
            #                 inputs={"x": begin.name},
            #                 outputs=[reshape_name],
            #                 shape=shape)
            #             inputs['offsets'] = reshape_name
            begin = self.decoder.infer_tensor(begin).tolist()
            attrs['offsets'] = begin
881
        if size.layer_type == "Const":
882
            size = size.value.tolist()
J
jiangjiajun 已提交
883
            attrs['shape'] = size
884
        else:
885
            shape = size.out_shapes[0]
J
jiangjiajun 已提交
886 887 888 889 890 891 892 893 894 895 896 897
            reshape_name = gen_name("slice", "reshape")
            program.add_layer(
                kernel="fluid.layers.reshape",
                inputs={"x": size.name},
                outputs=[reshape_name],
                shape=shape)
            inputs['shape'] = reshape_name
        program.add_layer(
            kernel="fluid.layers.crop_tensor",
            inputs=inputs,
            outputs=[node.name],
            **attrs)
898

J
jiangjiajun 已提交
899 900 901 902 903 904
    def ResizeNearestNeighbor(self, node):
        input = self.graph.get_node(node.layer.input[0])
        resize_shape = self.graph.get_node(node.layer.input[1])
        data_format = "NHWC"
        inputs = {"input": input.name}
        attrs = {"align_corners": node.get_attr("align_corners")}
905

J
jiangjiajun 已提交
906 907 908 909 910 911 912 913 914 915 916 917
        if resize_shape.layer_type == "Const":
            resize_shape = resize_shape.value.tolist()
            attrs["out_shape"] = resize_shape
        else:
            shape = resize_shape.out_shapes[0]
            reshape_name = gen_name("resize_nearest", "reshape")
            program.add_layer(
                kernel="fluid.layers.reshape",
                inputs={"x": resize_shape.name},
                outputs=[reshape_name],
                shape=shape)
            inputs["out_shape"] = reshape_name
918

J
jiangjiajun 已提交
919 920 921 922 923 924 925 926 927 928 929 930 931 932
        if data_format == "NHWC":
            transpose_name = gen_name("resize_nearest", "reshape")
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": input.name},
                outputs=[transpose_name],
                perm=[0, 3, 1, 2])
            inputs["input"] = transpose_name

        program.add_layer(
            kernel="fluid.layers.resize_nearest",
            inputs=inputs,
            outputs=[node.name],
            **attrs)
933

J
jiangjiajun 已提交
934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950
        if data_format == "NHWC":
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": node.name},
                outputs=[node.name],
                perm=[0, 2, 3, 1])

    def ResizeBilinear(self, node):
        input = self.graph.get_node(node.layer.input[0])
        resize_shape = self.graph.get_node(node.layer.input[1])
        data_format = "NHWC"
        inputs = {"input": input.name}
        attrs = {"align_corners": node.get_attr("align_corners")}

        if resize_shape.layer_type == "Const":
            resize_shape = resize_shape.value.tolist()
            attrs["out_shape"] = resize_shape
951
        else:
J
jiangjiajun 已提交
952 953 954 955 956 957 958 959
            shape = resize_shape.out_shapes[0]
            reshape_name = gen_name("resize_bilinear", "reshape")
            program.add_layer(
                kernel="fluid.layers.reshape",
                inputs={"x": resize_shape.name},
                outputs=[reshape_name],
                shape=shape)
            inputs["out_shape"] = reshape_name
960

J
jiangjiajun 已提交
961 962 963 964 965 966 967 968 969 970 971 972 973 974
        if data_format == "NHWC":
            transpose_name = gen_name("resize_bilinear", "reshape")
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": input.name},
                outputs=[transpose_name],
                perm=[0, 3, 1, 2])
            inputs["input"] = transpose_name

        program.add_layer(
            kernel="fluid.layers.resize_bilinear",
            inputs=inputs,
            outputs=[node.name],
            **attrs)
975

J
jiangjiajun 已提交
976 977 978 979 980 981
        if data_format == "NHWC":
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": node.name},
                outputs=[node.name],
                perm=[0, 2, 3, 1])
982

J
jiangjiajun 已提交
983 984 985 986 987 988 989 990
    def Cast(self, node):
        input = self.graph.get_node(node.layer.input[0])
        dtype = node.dtype
        program.add_layer(
            kernel="fluid.layers.cast",
            inputs={"x": input.name},
            outputs=[node.name],
            dtype=string(dtype))
991

J
jiangjiajun 已提交
992 993 994
    def Sum(self, node):
        input = self.graph.get_node(node.layer.input[0])
        reduce_idx = self.graph.get_node(node.layer.input[1])
995 996 997 998
        assert reduce_idx.layer_type == "Const", "Only support Const parameter[reduce_idx]"
        keep_dims = node.get_attr("keep_dims")
        dim = reduce_idx.value.tolist()

J
jiangjiajun 已提交
999 1000 1001 1002 1003 1004
        program.add_layer(
            kernel="fluid.layers.reduce_sum",
            inputs={"input": input.name},
            outputs=[node.name],
            dim=dim,
            keep_dim=keep_dims)
1005

J
jiangjiajun 已提交
1006 1007 1008
    def Max(self, node):
        input = self.graph.get_node(node.layer.input[0])
        reduce_idx = self.graph.get_node(node.layer.input[1])
1009 1010 1011
        assert reduce_idx.layer_type == "Const", "Only support Const parameter[reduce_idx]"
        keep_dims = node.get_attr("keep_dims")
        dim = reduce_idx.value.tolist()
J
jiangjiajun 已提交
1012 1013 1014 1015 1016 1017
        program.add_layer(
            kernel="fluid.layers.reduce_max",
            inputs={"input": input.name},
            outputs=[node.name],
            dim=dim,
            keep_dim=keep_dims)
1018 1019

    def RandomUniform(self, node):
M
update  
mamingjie-China 已提交
1020
        shape = self.graph.get_node(node.layer.input[0])
1021 1022
        if shape.layer_type == "Const":
            shape = shape.value.tolist()
J
jiangjiajun 已提交
1023 1024 1025 1026 1027 1028 1029
            program.add_layer(
                kernel="fluid.layers.uniform_random",
                inputs={},
                outputs=[node.name],
                shape=shape,
                min=0.0,
                max=0.9999)
1030
        else:
J
jiangjiajun 已提交
1031 1032 1033 1034 1035 1036
            program.add_layer(
                kernel="fluid.layers.uniform_random",
                inputs={'shape': shape.name},
                outputs=[node.name],
                min=0.0,
                max=0.9999)
M
mamingjie-China 已提交
1037

J
jiangjiajun 已提交
1038 1039 1040 1041 1042 1043 1044 1045 1046
    def Conv2DBackpropInput(self, node):
        out_shape = self.graph.get_node(node.layer.input[0])
        kernel = self.graph.get_node(node.layer.input[1])
        input = self.graph.get_node(node.layer.input[2])

        assert kernel.layer_type == "Const", "Kernel of Conv2DBackpropInput should be Const"

        if out_shape.layer_type == "Const":
            out_shape = out_shape.value.tolist()
M
mamingjie-China 已提交
1047
        else:
J
jiangjiajun 已提交
1048 1049
            out_shape = self.decoder.infer_shape_tensor(out_shape,
                                                        node.out_shapes[0])
M
mamingjie-China 已提交
1050

J
jiangjiajun 已提交
1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061
        in_shape = input.out_shapes[0]
        if in_shape.count(-1) > 2:
            in_shape = self.decoder.infer_tensor(input).shape
        k_size = kernel.out_shapes[0]
        if k_size.count(-1) > 2:
            k_size = self.decoder.infer_tensor(kernel).shape

        pad_mode = node.get_attr("padding").decode()
        strides = node.get_attr("strides")
        dilations = node.get_attr("dilations")
        data_format = node.get_attr("data_format").decode()
M
mamingjie-China 已提交
1062

J
jiangjiajun 已提交
1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097
        program.parameters[kernel.layer_name.replace(
            '/', '_')] = numpy.transpose(kernel.value, (3, 2, 0, 1))

        input_name = input.name
        if data_format == "NHWC":
            in_shape = [in_shape[i] for i in [0, 3, 1, 2]]
            strides = [strides[i] for i in [0, 3, 1, 2]]
            dilations = [dilations[i] for i in [0, 3, 1, 2]]
            transpose_name = gen_name("conv2dbackpropinput", "transpose")
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": input.name},
                outputs=[transpose_name],
                perm=[0, 3, 1, 2])
            input_name = transpose_name

        program.add_layer(
            kernel="fluid.layers.conv2d_transpose",
            inputs={"input": input_name},
            outputs=[node.name],
            bias_attr=False,
            param_attr=string(kernel.layer_name),
            num_filters=k_size[2],
            filter_size=k_size[0:2],
            stride=strides[2:4],
            dilation=dilations[2:4],
            padding=string(pad_mode),
            output_size=out_shape[1:3])

        if data_format == "NHWC":
            program.add_layer(
                kernel="fluid.layers.transpose",
                inputs={"x": node.name},
                outputs=[node.name],
                perm=[0, 2, 3, 1])
M
mamingjie-China 已提交
1098 1099

    def Tile(self, node):
M
update  
mamingjie-China 已提交
1100 1101
        input = self.graph.get_node(node.layer.input[0])
        expand_times = self.graph.get_node(node.layer.input[1])
M
mamingjie-China 已提交
1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116
        inputs = {"x": input.name}
        attr = dict()
        if expand_times.layer_type == "Const":
            expand_times = expand_times.value.tolist()
            attr["expand_times"] = expand_times
        else:
            inputs["expand_times"] = expand_times.name

        program.add_layer(
            kernel="fluid.layers.expand",
            inputs=inputs,
            outputs=[node.name],
            **attr)

    def Range(self, node):
M
update  
mamingjie-China 已提交
1117 1118 1119
        start = self.graph.get_node(node.layer.input[0])
        limit = self.graph.get_node(node.layer.input[1])
        delta = self.graph.get_node(node.layer.input[2])
M
mamingjie-China 已提交
1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141
        inputs = dict()
        attr = dict()

        if start.layer_type == "Const":
            attr["start"] = start.value
        else:
            inputs["start"] = start.name
        if limit.layer_type == "Const":
            attr["end"] = limit.value
        else:
            inputs["end"] = limit.name
        if delta.layer_type == "Const":
            attr["step"] = delta.value
        else:
            inputs["step"] = delta.name
        attr["dtype"] = string(node.dtype)

        program.add_layer(
            kernel="fluid.layers.range",
            inputs=inputs,
            outputs=[node.name],
            **attr)
M
update  
mamingjie-China 已提交
1142 1143

    def SquaredDifference(self, node):
M
update  
mamingjie-China 已提交
1144 1145
        x = self.graph.get_node(node.layer.input[0])
        y = self.graph.get_node(node.layer.input[1])
M
update  
mamingjie-China 已提交
1146 1147 1148 1149 1150 1151
        inputs = {"x": x.name, "y": y.name}
        program.add_layer(
            "fluid.layers.elementwise_sub", inputs=inputs, outputs=[node.name])
        inputs = {"x": node.name, "y": node.name}
        program.add_layer(
            "fluid.layers.elementwise_mul", inputs=inputs, outputs=[node.name])
M
mamingjie-China 已提交
1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201

    def OneHot(self, node):
        input = self.graph.get_node(node.layer.input[0])
        depth = self.graph.get_node(node.layer.input[1])
        on_value = self.graph.get_node(node.layer.input[2])
        off_value = self.graph.get_node(node.layer.input[3])
        assert depth.layer_type == 'Const', 'Parameter depth should be Const in OneHot'
        assert on_value.layer_type == 'Const', 'Parameter on_value should be Const in OneHot'
        assert off_value.layer_type == 'Const', 'Parameter off_value should be Const in OneHot'

        attr = {'depth': depth.value}
        on_value = on_value.value
        off_value = off_value.value
        assert math.fabs(on_value -
                         1.0) < 1e-06, "on_value should be 1 in OneHot"
        assert math.fabs(off_value -
                         0.0) < 1e-06, "off_value should be 0 in OneHot"

        program.add_layer(
            "fluid.one_hot",
            inputs={"input": input.name},
            outputs=[node.name],
            depth=depth.value)

    def Pow(self, node):
        x = self.graph.get_node(node.layer.input[0])
        factor = self.graph.get_node(node.layer.input[1])
        inputs = {"x": x.name}
        attr = dict()
        if factor.layer_type == 'Const':
            attr["factor"] = factor.value.tolist()
        else:
            inputs["factor"] = factor.name
        program.add_layer(
            "fluid.layers.pow", inputs=inputs, outputs=[node.name], **attr)

    def All(self, node):
        input = self.graph.get_node(node.layer.input[0])
        reduce_idx = self.graph.get_node(node.layer.input[1])
        assert reduce_idx.layer_type == "Const", "Only support Const parameter[reduce_idx]"
        attr = dict()
        attr["dim"] = reduce_idx.value.tolist()
        attr["keep_dim"] = node.get_attr("keep_dims")

        program.add_layer(
            "fluid.layers.reduce_all",
            inputs={"input": input.name},
            outputs=[node.name],
            **attr)

J
jiangjiajun 已提交
1202 1203
        node.layer.attr['dtype'].type = 10

M
mamingjie-China 已提交
1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225
    def GatherV2(self, node):
        embeddings = self.graph.get_node(node.layer.input[0])
        index = self.graph.get_node(node.layer.input[1])
        axis = self.graph.get_node(node.layer.input[2])
        assert axis.layer_type == 'Const', "Only support Const parameter[axis]"
        axis = axis.value.tolist()
        assert axis == 0, "Only support axis=0 in GatherV2 OP"
        index_name = index.name
        if len(index.out_shapes[0]) != 1:
            reshape_name = gen_name("gather", "reshape")
            index_name = reshape_name
            program.add_layer(
                "fluid.layers.reshape",
                inputs={"x": index.name},
                outputs=[reshape_name],
                shape=[-1])
        inputs = {'input': embeddings.name, 'index': index_name}
        program.add_layer(
            "fluid.layers.gather",
            inputs=inputs,
            outputs=[node.name],
            overwrite=False)
J
jiangjiajun 已提交
1226 1227 1228 1229 1230 1231 1232
        if len(index.out_shapes[0]) != 1:
            out_shape = node.out_shapes[0]
            program.add_layer(
                kernel="fluid.layers.reshape",
                inputs={"x": node.name},
                outputs=[node.name],
                shape=out_shape)
M
mamingjie-China 已提交
1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250

    def ExpandDims(self, node):
        x = self.graph.get_node(node.layer.input[0], copy=True)
        y = self.graph.get_node(node.layer.input[1], copy=True)
        inputs = {"input": x.name}
        attr = dict()
        if y.layer_type == 'Const':
            dim = y.value.tolist()
            if not isinstance(dim, list):
                dim = [dim]
            attr['axes'] = dim
        else:
            inputs['axes'] = y.name
        program.add_layer(
            "fluid.layers.unsqueeze",
            inputs=inputs,
            outputs=[node.name],
            **attr)