PPLCNet/base_layer.py at master · Hiwyl/PPLCNet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

from paddle import nn
import re


class Identity(nn.Layer):
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, inputs):
        return inputs


class TheseusLayer(nn.Layer):
    def __init__(self, *args, **kwargs):
        super(TheseusLayer, self).__init__()
        self.res_dict = {}

    # stop doesn't work when stop layer has a parallel branch.
    def stop_after(self, stop_layer_name: str):
        after_stop = False
        for layer_i in self._sub_layers:
            if after_stop:
                self._sub_layers[layer_i] = Identity()
                continue
            layer_name = self._sub_layers[layer_i].full_name()
            if layer_name == stop_layer_name:
                after_stop = True
                continue
            if isinstance(self._sub_layers[layer_i], TheseusLayer):
                after_stop = self._sub_layers[layer_i].stop_after(
                    stop_layer_name)
        return after_stop

    def update_res(self, return_patterns):
        if not return_patterns or isinstance(self, WrapLayer):
            return
        for layer_i in self._sub_layers:
            layer_name = self._sub_layers[layer_i].full_name()
            if isinstance(self._sub_layers[layer_i], (nn.Sequential, nn.LayerList)):
                self._sub_layers[layer_i] = wrap_theseus(self._sub_layers[layer_i], self.res_dict)
                self._sub_layers[layer_i].update_res(return_patterns)
            else:
                for return_pattern in return_patterns:
                    if re.match(return_pattern, layer_name):
                        if not isinstance(self._sub_layers[layer_i], TheseusLayer):
                            self._sub_layers[layer_i] = wrap_theseus(self._sub_layers[layer_i], self.res_dict)
                        else:
                            self._sub_layers[layer_i].res_dict = self.res_dict

                        self._sub_layers[layer_i].register_forward_post_hook(
                            self._sub_layers[layer_i]._save_sub_res_hook)
            if isinstance(self._sub_layers[layer_i], TheseusLayer):
                self._sub_layers[layer_i].res_dict = self.res_dict
                self._sub_layers[layer_i].update_res(return_patterns)

    def _save_sub_res_hook(self, layer, input, output):
        self.res_dict[layer.full_name()] = output

    def _return_dict_hook(self, layer, input, output):
        res_dict = {"output": output}
        for res_key in list(self.res_dict):
            res_dict[res_key] = self.res_dict.pop(res_key)
        return res_dict

    def replace_sub(self, layer_name_pattern, replace_function, recursive=True):
        for layer_i in self._sub_layers:
            layer_name = self._sub_layers[layer_i].full_name()
            if re.match(layer_name_pattern, layer_name):
                self._sub_layers[layer_i] = replace_function(self._sub_layers[layer_i])
            if recursive:
                if isinstance(self._sub_layers[layer_i], TheseusLayer):
                    self._sub_layers[layer_i].replace_sub(
                        layer_name_pattern, replace_function, recursive)
                elif isinstance(self._sub_layers[layer_i], (nn.Sequential, nn.LayerList)):
                    for layer_j in self._sub_layers[layer_i]._sub_layers:
                        self._sub_layers[layer_i]._sub_layers[layer_j].replace_sub(
                            layer_name_pattern, replace_function, recursive)

    '''
    example of replace function:
    def replace_conv(origin_conv: nn.Conv2D):
        new_conv = nn.Conv2D(
            in_channels=origin_conv._in_channels,
            out_channels=origin_conv._out_channels,
            kernel_size=origin_conv._kernel_size,
            stride=2
        )
        return new_conv
        '''


class WrapLayer(TheseusLayer):
    def __init__(self, sub_layer, res_dict=None):
        super(WrapLayer, self).__init__()
        self.sub_layer = sub_layer
        self.name = sub_layer.full_name()
        if res_dict is not None:
            self.res_dict = res_dict

    def full_name(self):
        return self.name

    def forward(self, *inputs, **kwargs):
        return self.sub_layer(*inputs, **kwargs)

    def update_res(self, return_patterns):
        if not return_patterns or not isinstance(self.sub_layer, (nn.Sequential, nn.LayerList)):
            return
        for layer_i in self.sub_layer._sub_layers:
            if isinstance(self.sub_layer._sub_layers[layer_i], (nn.Sequential, nn.LayerList)):
                self.sub_layer._sub_layers[layer_i] = wrap_theseus(self.sub_layer._sub_layers[layer_i], self.res_dict)
                self.sub_layer._sub_layers[layer_i].update_res(return_patterns)
            elif isinstance(self.sub_layer._sub_layers[layer_i], TheseusLayer):
                self.sub_layer._sub_layers[layer_i].res_dict = self.res_dict

            layer_name = self.sub_layer._sub_layers[layer_i].full_name()
            for return_pattern in return_patterns:
                if re.match(return_pattern, layer_name):
                    self.sub_layer._sub_layers[layer_i].register_forward_post_hook(
                        self._sub_layers[layer_i]._save_sub_res_hook)

            if isinstance(self.sub_layer._sub_layers[layer_i], TheseusLayer):
                self.sub_layer._sub_layers[layer_i].update_res(return_patterns)


def wrap_theseus(sub_layer, res_dict=None):
    wrapped_layer = WrapLayer(sub_layer, res_dict)
    return wrapped_layer