interpolate_bilinear_fuser.py

#   Copyright (c) 2020  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.

import copy
import numpy as np
from x2paddle.optimizer.pattern_matcher import FuseBase
from x2paddle.core.program import PaddleGraph, PaddleLayer
from x2paddle.core.util import *


class DygraphInterpolateBilinearFuser(FuseBase):
    def __init__(self):
        super(DygraphInterpolateBilinearFuser, self).__init__(graph_type="dygraph")
        import torch
        torch_version = torch.__version__
        torch_version_part = torch_version.split(".")
        if int(torch_version_part[0]) == 1 and int(torch_version_part[1]) > 5:
            self.version_gt_150 = True
        else:
            self.version_gt_150 = False

    def build_pattern(self):
        """ 描述需要替换的双线性插值图结构。
        interpolate_bilinear层模式python实现代码示例:
            x2195 = x2181.shape
            x2195 = len(x2195)
            x2196 = x2195 - 2
            x2197 = []
            for _x2199 in range(x2196):
                x2197.append(None)
            x2200 = (x2181, x8, None, None)
            ...
            x2267 = x2266 == 3
            if x2267 :
                raise RaiseException('Exception')
                x2268 = None
            else:
                x2270 = x2181.shape
                x2270 = len(x2270)
                x2271 = x2270 == 4
                if x2271 :
                    x2274 = x2197[0]
                    x2275 = x2197[1]
                    x2233_isinstance = isinstance(x2233, paddle.fluid.Variable)
                    if x2233_isinstance :
                        x2233 = x2233.numpy().tolist()
                    x2276 = paddle.nn.functional.interpolate(x=x2181, size=x2233, scale_factor=x2274, align_corners=False, align_mode=0, mode='bilinear')
                    x2272 = x2276
                else:
                    x2277 = x2181.shape
                    x2277 = len(x2277)
                    x2278 = x2277 == 5
                    if x2278 :
                        raise RaiseException('Exception')
                    else:
                        raise RaiseException('Exception')
                    x2272 = None
                x2268 = x2272
        """

        def gen_name(id):
            return "x" + str(id)

        if self.version_gt_150:
            self.pattern.add_layer(
                "prim.shape",
                inputs={"input": "interpolate-input-0"},
                outputs=[gen_name(9)])
            self.pattern.add_layer(
                "prim.len",
                inputs={"input": gen_name(9)},
                outputs=[gen_name(9)])
            self.pattern.add_layer(
                "prim.sub",
                inputs={"x": gen_name(9)},
                outputs=[gen_name(10)],
                y=2)
            self.pattern.add_layer(
                "prim.list", inputs={}, outputs=[gen_name(11)])
            self.pattern.add_layer(
                "prim.loop",
                inputs={"input": gen_name(10)},
                outputs=[gen_name(12.1), gen_name(12.2)])
            loop_layer = self.pattern.layers[list(self.pattern.layers.keys())[
                -1]]
            pattern_block = PaddleGraph(loop_layer, graph_type="dygraph")
            pattern_block.add_layer(
                "prim.append",
                inputs={"list": gen_name(11)},
                outputs=[],
                element=None)
            loop_layer.inputs["input-0"] = gen_name(11)
            loop_layer.add_block(pattern_block)
            self.pattern.add_layer(
                "prim.tuple",
                inputs={
                    "input0": "interpolate-input-0",
                    "input1": "interpolate-input-4",
                },
                outputs=[gen_name(12)],
                input2=None,
                input3=None)
            
            self.pattern.add_layer(
                "prim.eq",
                inputs={"x": "interpolate-input-2"},
                outputs=[gen_name(10.1)],
                y=3)
            
            self.pattern.add_layer(
                "prim.if",
                inputs={"input": gen_name(10.1)},
                outputs=[gen_name(14)])
            if_layer1 = self.pattern.layers[list(self.pattern.layers.keys())[
                -1]]
            pattern_block = PaddleGraph(parent_layer=if_layer1, graph_type="dygraph")
            pattern_block.add_layer(
                "prim.exception",
                inputs={},
                outputs=[gen_name(15)],
                input="Exception")
            pattern_block.add_layer(
                "prim.equal", inputs={}, outputs=[gen_name(14)], input=None)
            if_layer1.add_block(pattern_block)
            pattern_block = PaddleGraph(parent_layer=if_layer1, graph_type="dygraph")
            pattern_block.add_layer(
                "prim.shape",
                inputs={"input": "interpolate-input-0"},
                outputs=[gen_name(18)])
            pattern_block.add_layer(
                "prim.len",
                inputs={"input": gen_name(18)},
                outputs=[gen_name(18)])
            pattern_block.add_layer(
                "prim.eq",
                inputs={"x": gen_name(18)},
                outputs=[gen_name(19)],
                y=4)
            
            
            pattern_block.add_layer(
                "prim.if",
                inputs={"input": gen_name(19)},
                outputs=[gen_name(20)])
            if_layer2 = pattern_block.layers[list(pattern_block.layers.keys())[
                -1]]
            pattern_block_block = PaddleGraph(parent_layer=if_layer2, graph_type="dygraph")            
            pattern_block_block.add_layer(
                "prim.getitem",
                inputs={"list": gen_name(11)},
                outputs=[gen_name(21)],
                element=0)
            pattern_block_block.add_layer(
                "prim.getitem",
                inputs={"list": gen_name(11)},
                outputs=[gen_name(22)],
                element=1)
            pattern_block_block.add_layer(
                "prim.isinstance",
                inputs={"input": "interpolate-input-3"},
                outputs=["interpolate-input-0_isinstance"],
                cls="paddle.fluid.Variable")
            pattern_block_block.add_layer(
                "prim.if", {"input": "interpolate-input-0_isinstance"},
                outputs=["interpolate-input-0_if1"])
            if_layer_isinstance = pattern_block_block.layers[list(
                pattern_block_block.layers.keys())[-1]]
            pattern_block_block_block = PaddleGraph(
                if_layer_isinstance, graph_type="dygraph")
            pattern_block_block_block.add_layer(
                "prim.var2list",
                inputs={"input": "interpolate-input-3"},
                outputs=["interpolate-input-3"])
            if_layer_isinstance.add_block(pattern_block_block_block)
            pattern_block_block_block = PaddleGraph(
                if_layer_isinstance, graph_type="dygraph")
            if_layer_isinstance.add_block(pattern_block_block_block)
            if_layer_isinstance.inputs["input-0"] = "interpolate-input-3"
            pattern_block_block.add_layer(
                "paddle.nn.functional.interpolate",
                inputs={
                    "input": "interpolate-input-0",
                    "size": "interpolate-input-3",
                    "scale_factor": gen_name(21)
                },
                outputs=[gen_name(23)])
            pattern_block_block.add_layer(
                "prim.equal",
                inputs={"input": gen_name(23)},
                outputs=[gen_name(20)])
            if_layer2.add_block(pattern_block_block)
            pattern_block_block = PaddleGraph(if_layer2, graph_type="dygraph")
            pattern_block_block.add_layer(
                "prim.shape",
                inputs={"input": "interpolate-input-0"},
                outputs=[gen_name(24)])
            pattern_block_block.add_layer(
                "prim.len",
                inputs={"input": gen_name(24)},
                outputs=[gen_name(24)])
            pattern_block_block.add_layer(
                "prim.eq",
                inputs={"x": gen_name(24)},
                outputs=[gen_name(25)],
                y=5)
            pattern_block_block.add_layer(
                "prim.if",
                inputs={"input": gen_name(25)},
                outputs=[gen_name(26)])
            if_layer3 = pattern_block_block.layers[list(
                pattern_block_block.layers.keys())[-1]]
            pattern_block_block_block = PaddleGraph(
                parent_layer=if_layer3, graph_type="dygraph")
            pattern_block_block_block.add_layer(
                "prim.exception",
                inputs={},
                outputs=[gen_name(27)],
                input="Exception")
            if_layer3.add_block(pattern_block_block_block)
            pattern_block_block_block = PaddleGraph(
                parent_layer=if_layer3, graph_type="dygraph")
            pattern_block_block_block.add_layer(
                "prim.exception",
                inputs={},
                outputs=[gen_name(28)],
                input="Exception")
            if_layer3.add_block(pattern_block_block_block)
            pattern_block_block.add_layer(
                "prim.equal", inputs={}, outputs=[gen_name(20)], input=None)
            if_layer2.add_block(pattern_block_block)
            if_layer2.inputs.update({
                "input-0": "interpolate-input-0",
                "input-1": "interpolate-input-3",
                "input-2": "interpolate-input-3",
                "input-3": gen_name(11),
                "input-5": gen_name(11),
            })
            pattern_block.add_layer(
                "prim.equal",
                inputs={"input": gen_name(20)},
                outputs=[gen_name(14)])
            if_layer1.add_block(pattern_block)
            if_layer1.inputs.update({
                'input-2': 'interpolate-input-0', 
                'input-4': gen_name(11), 
                'input-6': gen_name(11), 
                'input-8': 'interpolate-input-0', 
                'input-9': 'interpolate-input-3', 
                'input-10': 'interpolate-input-0'
            })
            self.pattern.build(inputs={
                "input-0": "interpolate-input-0",
                "input-1": "interpolate-input-1",
                "input-2": "interpolate-input-2",
                "input-3": "interpolate-input-3",
                "input-4": "interpolate-input-4"
            })
            
            
            

    def insert_new_layer(self, graph, parameters, matches):
        new_layer = self.gen_new_layer(parameters, matches)
        global_layers = graph.get_global_layers()
        new_matches = dict()
        is_match = False
        for layer_id, layer in global_layers.items():
            if layer_id == list(matches.keys())[0] and not is_match:
                new_matches[layer_id] = layer
                is_match = True
            if is_match:
                new_matches[layer_id] = layer
                if layer_id == list(matches.keys())[-1]:
                    break
        new_layer_id = new_layer.layer_id
        graph.layers[new_layer_id] = new_layer
        new_matches.pop(new_layer_id)
        matches.clear()
        for layer_id, layer in new_matches.items():
            matches[layer_id] = layer
        

    def gen_new_layer(self, parameters, matches):
        layers = list()
        layers_id = list(matches.keys())
        layer = matches[layers_id[6]]
        size = layer.inputs["input1"]
        layer = matches[layers_id[19]]
        new_layer = copy.deepcopy(layer)
        layer = matches[layers_id[9]]
        new_layer.outputs[0] = layer.outputs[0]
        new_layer.layer_id = layers_id[7]
        new_layer.inputs.pop("scale_factor")
        new_layer.inputs["size"] = size
        return new_layer