keras-cv/keras_cv/layers/efficient_multihead_attention.py at 314dc6b946977b60955a369a6de582498864d1c9 · DavidLandup0/keras-cv · 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
from keras_cv.backend import keras

"""
Based on: https://github.com/sithu31296/semantic-segmentation/blob/main/semseg/models/backbones/mit.py
"""


@keras.saving.register_keras_serializable(package="keras_cv")
class EfficientMultiheadAttention(keras.layers.Layer):
    def __init__(self, project_dim, num_heads, sr_ratio):
        super().__init__()
        self.num_heads = num_heads
        self.sr_ratio = sr_ratio
        self.scale = (project_dim // num_heads) ** -0.5
        self.q = keras.layers.Dense(project_dim)
        self.k = keras.layers.Dense(project_dim)
        self.v = keras.layers.Dense(project_dim)
        self.proj = keras.layers.Dense(project_dim)

        if sr_ratio > 1:
            self.sr = keras.layers.Conv2D(
                filters=project_dim,
                kernel_size=sr_ratio,
                strides=sr_ratio,
                padding="same",
            )
            self.norm = keras.layers.LayerNormalization()

    def call(self, x, H, W):
        input_shape = x.shape

        q = self.q(x)
        q = keras.ops.reshape.reshape(
            q,
            (
                input_shape[0],
                input_shape[1],
                self.num_heads,
                input_shape[2] // self.num_heads,
            ),
        )

        q = q.transpose([0, 2, 1, 3])

        if self.sr_ratio > 1:
            x = keras.ops.reshape(
                keras.ops.transpose(x, [0, 2, 1]),
                (input_shape[0], H, W, input_shape[2]),
            )
            x = self.sr(x)
            x = keras.ops.reshape(x, [input_shape[0], input_shape[2], -1])
            x = keras.ops.transpose(x, [0, 2, 1])
            x = self.norm(x)

        k = self.k(x)
        v = self.v(x)

        k = keras.ops.reshape(
            keras.ops.transpose(k, [0, 2, 1, 3]),
            [
                input_shape[0],
                -1,
                self.num_heads,
                input_shape[2] // self.num_heads,
            ],
        )

        v = keras.ops.reshape(
            keras.ops.transpose(v, [0, 2, 1, 3]),
            [
                input_shape[0],
                -1,
                self.num_heads,
                input_shape[2] // self.num_heads,
            ],
        )

        attn = (q @ keras.ops.transpose(x, [0, 1, 3, 2])) * self.scale
        attn = keras.nn.ops.softmax(attn, axis=-1)

        attn = attn @ v
        attn = keras.ops.reshape(
            keras.ops.transpose(attn, [0, 2, 1, 3]),
            [input_shape[0], input_shape[1], input_shape[2]],
        )
        x = self.proj(attn)
        return x