attention2d.py 10.6 KB
Newer Older
1
2
3
4
5
6
import math

import torch
from torch import nn


Patrick von Platen's avatar
Patrick von Platen committed
7
# unet_grad_tts.py
Patrick von Platen's avatar
Patrick von Platen committed
8
# TODO(Patrick) - weird linear attention layer. Check with: https://github.com/huawei-noah/Speech-Backbones/issues/15
Patrick von Platen's avatar
Patrick von Platen committed
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class LinearAttention(torch.nn.Module):
    def __init__(self, dim, heads=4, dim_head=32):
        super(LinearAttention, self).__init__()
        self.heads = heads
        self.dim_head = dim_head
        hidden_dim = dim_head * heads
        self.to_qkv = torch.nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
        self.to_out = torch.nn.Conv2d(hidden_dim, dim, 1)

    def forward(self, x):
        b, c, h, w = x.shape
        qkv = self.to_qkv(x)
        q, k, v = (
            qkv.reshape(b, 3, self.heads, self.dim_head, h, w)
            .permute(1, 0, 2, 3, 4, 5)
            .reshape(3, b, self.heads, self.dim_head, -1)
        )
        k = k.softmax(dim=-1)
        context = torch.einsum("bhdn,bhen->bhde", k, v)
        out = torch.einsum("bhde,bhdn->bhen", context, q)
        out = out.reshape(b, self.heads, self.dim_head, h, w).reshape(b, self.heads * self.dim_head, h, w)
        return self.to_out(out)

32

Patrick von Platen's avatar
Patrick von Platen committed
33
# the main attention block that is used for all models
Patrick von Platen's avatar
Patrick von Platen committed
34
35
36
37
38
39
40
41
42
43
44
45
46
class AttentionBlock(nn.Module):
    """
    An attention block that allows spatial positions to attend to each other.

    Originally ported from here, but adapted to the N-d case.
    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
    """

    def __init__(
        self,
        channels,
        num_heads=1,
        num_head_channels=-1,
Patrick von Platen's avatar
Patrick von Platen committed
47
        num_groups=32,
Patrick von Platen's avatar
Patrick von Platen committed
48
49
        use_checkpoint=False,
        encoder_channels=None,
50
        use_new_attention_order=False,  # TODO(Patrick) -> is never used, maybe delete?
Patrick von Platen's avatar
Patrick von Platen committed
51
        overwrite_qkv=False,
Patrick von Platen's avatar
Patrick von Platen committed
52
53
        overwrite_linear=False,
        rescale_output_factor=1.0,
Patrick von Platen's avatar
Patrick von Platen committed
54
55
56
57
58
59
60
61
62
63
    ):
        super().__init__()
        self.channels = channels
        if num_head_channels == -1:
            self.num_heads = num_heads
        else:
            assert (
                channels % num_head_channels == 0
            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
            self.num_heads = channels // num_head_channels
Patrick von Platen's avatar
Patrick von Platen committed
64

Patrick von Platen's avatar
Patrick von Platen committed
65
        self.use_checkpoint = use_checkpoint
Patrick von Platen's avatar
Patrick von Platen committed
66
67
        self.norm = nn.GroupNorm(num_channels=channels, num_groups=num_groups, eps=1e-5, affine=True)
        self.qkv = nn.Conv1d(channels, channels * 3, 1)
Patrick von Platen's avatar
Patrick von Platen committed
68
        self.n_heads = self.num_heads
Patrick von Platen's avatar
Patrick von Platen committed
69
        self.rescale_output_factor = rescale_output_factor
Patrick von Platen's avatar
Patrick von Platen committed
70
71

        if encoder_channels is not None:
Patrick von Platen's avatar
Patrick von Platen committed
72
            self.encoder_kv = nn.Conv1d(encoder_channels, channels * 2, 1)
Patrick von Platen's avatar
Patrick von Platen committed
73

Patrick von Platen's avatar
Patrick von Platen committed
74
        self.proj_out = zero_module(nn.Conv1d(channels, channels, 1))
Patrick von Platen's avatar
Patrick von Platen committed
75

Patrick von Platen's avatar
Patrick von Platen committed
76
77
78
        self.overwrite_qkv = overwrite_qkv
        if overwrite_qkv:
            in_channels = channels
Patrick von Platen's avatar
Patrick von Platen committed
79
            self.norm = nn.GroupNorm(num_channels=channels, num_groups=num_groups, eps=1e-6)
Patrick von Platen's avatar
Patrick von Platen committed
80
81
82
83
            self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
            self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
            self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
            self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
Patrick von Platen's avatar
Patrick von Platen committed
84

Patrick von Platen's avatar
Patrick von Platen committed
85
86
87
88
89
90
91
92
93
        self.overwrite_linear = overwrite_linear
        if self.overwrite_linear:
            num_groups = min(channels // 4, 32)
            self.norm = nn.GroupNorm(num_channels=channels, num_groups=num_groups, eps=1e-6)
            self.NIN_0 = NIN(channels, channels)
            self.NIN_1 = NIN(channels, channels)
            self.NIN_2 = NIN(channels, channels)
            self.NIN_3 = NIN(channels, channels)

Patrick von Platen's avatar
Patrick von Platen committed
94
        self.is_overwritten = False
95

Patrick von Platen's avatar
Patrick von Platen committed
96
97
98
99
    def set_weights(self, module):
        if self.overwrite_qkv:
            qkv_weight = torch.cat([module.q.weight.data, module.k.weight.data, module.v.weight.data], dim=0)[:, :, :, 0]
            qkv_bias = torch.cat([module.q.bias.data, module.k.bias.data, module.v.bias.data], dim=0)
Patrick von Platen's avatar
Patrick von Platen committed
100

Patrick von Platen's avatar
Patrick von Platen committed
101
102
103
            self.qkv.weight.data = qkv_weight
            self.qkv.bias.data = qkv_bias

Patrick von Platen's avatar
Patrick von Platen committed
104
            proj_out = zero_module(nn.Conv1d(self.channels, self.channels, 1))
Patrick von Platen's avatar
Patrick von Platen committed
105
106
            proj_out.weight.data = module.proj_out.weight.data[:, :, :, 0]
            proj_out.bias.data = module.proj_out.bias.data
Patrick von Platen's avatar
Patrick von Platen committed
107

Patrick von Platen's avatar
Patrick von Platen committed
108
            self.proj_out = proj_out
Patrick von Platen's avatar
Patrick von Platen committed
109
110
111
112
113
114
        elif self.overwrite_linear:
            self.qkv.weight.data = torch.concat([self.NIN_0.W.data.T, self.NIN_1.W.data.T, self.NIN_2.W.data.T], dim=0)[:, :, None]
            self.qkv.bias.data = torch.concat([self.NIN_0.b.data, self.NIN_1.b.data, self.NIN_2.b.data], dim=0)

            self.proj_out.weight.data = self.NIN_3.W.data.T[:, :, None]
            self.proj_out.bias.data = self.NIN_3.b.data
Patrick von Platen's avatar
Patrick von Platen committed
115
116
117
118
119
120
121
122

    def forward(self, x, encoder_out=None):
        if self.overwrite_qkv and not self.is_overwritten:
            self.set_weights(self)
            self.is_overwritten = True

        b, c, *spatial = x.shape
        hid_states = self.norm(x).view(b, c, -1)
Patrick von Platen's avatar
Patrick von Platen committed
123

Patrick von Platen's avatar
Patrick von Platen committed
124
        qkv = self.qkv(hid_states)
Patrick von Platen's avatar
Patrick von Platen committed
125
126
127
128
        bs, width, length = qkv.shape
        assert width % (3 * self.n_heads) == 0
        ch = width // (3 * self.n_heads)
        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
Patrick von Platen's avatar
Patrick von Platen committed
129
130
131

        if encoder_out is not None:
            encoder_kv = self.encoder_kv(encoder_out)
Patrick von Platen's avatar
Patrick von Platen committed
132
133
134
135
            assert encoder_kv.shape[1] == self.n_heads * ch * 2
            ek, ev = encoder_kv.reshape(bs * self.n_heads, ch * 2, -1).split(ch, dim=1)
            k = torch.cat([ek, k], dim=-1)
            v = torch.cat([ev, v], dim=-1)
Patrick von Platen's avatar
Patrick von Platen committed
136

Patrick von Platen's avatar
Patrick von Platen committed
137
138
139
        scale = 1 / math.sqrt(math.sqrt(ch))
        weight = torch.einsum("bct,bcs->bts", q * scale, k * scale)  # More stable with f16 than dividing afterwards
        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
Patrick von Platen's avatar
Patrick von Platen committed
140

Patrick von Platen's avatar
Patrick von Platen committed
141
        a = torch.einsum("bts,bcs->bct", weight, v)
Patrick von Platen's avatar
Patrick von Platen committed
142
143
144
        h = a.reshape(bs, -1, length)

        h = self.proj_out(h)
Patrick von Platen's avatar
Patrick von Platen committed
145
        h = h.reshape(b, c, *spatial)
Patrick von Platen's avatar
Patrick von Platen committed
146

Patrick von Platen's avatar
Patrick von Platen committed
147
        result = x + h
Patrick von Platen's avatar
Patrick von Platen committed
148

Patrick von Platen's avatar
Patrick von Platen committed
149
        result = result / self.rescale_output_factor
Patrick von Platen's avatar
Patrick von Platen committed
150

Patrick von Platen's avatar
Patrick von Platen committed
151
        return result
Patrick von Platen's avatar
Patrick von Platen committed
152
153
154


# unet_score_estimation.py
Patrick von Platen's avatar
Patrick von Platen committed
155
#class AttnBlockpp(nn.Module):
156
157
#    """Channel-wise self-attention block. Modified from DDPM."""
#
Patrick von Platen's avatar
Patrick von Platen committed
158
159
160
161
162
163
164
165
166
167
168
169
170
#    def __init__(
#        self,
#        channels,
#        skip_rescale=False,
#        init_scale=0.0,
#        num_heads=1,
#        num_head_channels=-1,
#        use_checkpoint=False,
#        encoder_channels=None,
#        use_new_attention_order=False,  # TODO(Patrick) -> is never used, maybe delete?
#        overwrite_qkv=False,
#        overwrite_from_grad_tts=False,
#    ):
171
#        super().__init__()
Patrick von Platen's avatar
Patrick von Platen committed
172
173
#        num_groups = min(channels // 4, 32)
#        self.GroupNorm_0 = nn.GroupNorm(num_groups=num_groups, num_channels=channels, eps=1e-6)
174
175
176
177
178
179
#        self.NIN_0 = NIN(channels, channels)
#        self.NIN_1 = NIN(channels, channels)
#        self.NIN_2 = NIN(channels, channels)
#        self.NIN_3 = NIN(channels, channels, init_scale=init_scale)
#        self.skip_rescale = skip_rescale
#
Patrick von Platen's avatar
Patrick von Platen committed
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
#        self.channels = channels
#        if num_head_channels == -1:
#            self.num_heads = num_heads
#        else:
#            assert (
#                channels % num_head_channels == 0
#            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
#            self.num_heads = channels // num_head_channels
#
#        self.use_checkpoint = use_checkpoint
#        self.norm = normalization(channels, num_groups=num_groups, eps=1e-6, swish=None)
#        self.qkv = conv_nd(1, channels, channels * 3, 1)
#        self.n_heads = self.num_heads
#
#        if encoder_channels is not None:
#            self.encoder_kv = conv_nd(1, encoder_channels, channels * 2, 1)
#
#        self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
#
#        self.is_weight_set = False
#
#    def set_weights(self):
#        self.qkv.weight.data = torch.concat([self.NIN_0.W.data.T, self.NIN_1.W.data.T, self.NIN_2.W.data.T], dim=0)[:, :, None]
#        self.qkv.bias.data = torch.concat([self.NIN_0.b.data, self.NIN_1.b.data, self.NIN_2.b.data], dim=0)
#
#        self.proj_out.weight.data = self.NIN_3.W.data.T[:, :, None]
#        self.proj_out.bias.data = self.NIN_3.b.data
#
208
#    def forward(self, x):
Patrick von Platen's avatar
Patrick von Platen committed
209
210
211
212
#        if not self.is_weight_set:
#            self.set_weights()
#            self.is_weight_set = True
#
213
214
215
216
217
218
219
220
221
222
223
224
#        B, C, H, W = x.shape
#        h = self.GroupNorm_0(x)
#        q = self.NIN_0(h)
#        k = self.NIN_1(h)
#        v = self.NIN_2(h)
#
#        w = torch.einsum("bchw,bcij->bhwij", q, k) * (int(C) ** (-0.5))
#        w = torch.reshape(w, (B, H, W, H * W))
#        w = F.softmax(w, dim=-1)
#        w = torch.reshape(w, (B, H, W, H, W))
#        h = torch.einsum("bhwij,bcij->bchw", w, v)
#        h = self.NIN_3(h)
Patrick von Platen's avatar
Patrick von Platen committed
225
#
226
#        if not self.skip_rescale:
Patrick von Platen's avatar
Patrick von Platen committed
227
#            result = x + h
228
#        else:
Patrick von Platen's avatar
Patrick von Platen committed
229
230
231
232
233
234
235
236
237
238
239
240
241
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
275
276
277
278
279
#            result = (x + h) / np.sqrt(2.0)
#
#        result = self.forward_2(x)
#
#        return result
#
#    def forward_2(self, x, encoder_out=None):
#        b, c, *spatial = x.shape
#        hid_states = self.norm(x).view(b, c, -1)
#
#        qkv = self.qkv(hid_states)
#        bs, width, length = qkv.shape
#        assert width % (3 * self.n_heads) == 0
#        ch = width // (3 * self.n_heads)
#        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
#
#        if encoder_out is not None:
#            encoder_kv = self.encoder_kv(encoder_out)
#            assert encoder_kv.shape[1] == self.n_heads * ch * 2
#            ek, ev = encoder_kv.reshape(bs * self.n_heads, ch * 2, -1).split(ch, dim=1)
#            k = torch.cat([ek, k], dim=-1)
#            v = torch.cat([ev, v], dim=-1)
#
#        scale = 1 / math.sqrt(math.sqrt(ch))
#        weight = torch.einsum("bct,bcs->bts", q * scale, k * scale)  # More stable with f16 than dividing afterwards
#        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
#
#        a = torch.einsum("bts,bcs->bct", weight, v)
#        h = a.reshape(bs, -1, length)
#
#        h = self.proj_out(h)
#        h = h.reshape(b, c, *spatial)
#
#        return (x + h) / np.sqrt(2.0)

# TODO(Patrick) - this can and should be removed
def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


# TODO(Patrick) - remove once all weights have been converted -> not needed anymore then
class NIN(nn.Module):
    def __init__(self, in_dim, num_units, init_scale=0.1):
        super().__init__()
        self.W = nn.Parameter(torch.zeros(in_dim, num_units), requires_grad=True)
        self.b = nn.Parameter(torch.zeros(num_units), requires_grad=True)