වේගයෙන්ට්රාන්ස්ෆෝමර් බර

96import torch
97from torch import nn
98
99from labml_helpers.module import Module
100from labml_nn.transformers.feed_forward import FeedForward
101from labml_nn.transformers.mha import PrepareForMultiHeadAttention
102from labml_nn.utils import clone_module_list

105class DPFP(Module):

139    def __init__(self, nu: int = 1, eps: float = 1e-6):

144        super().__init__()
145        self.nu = nu
146        self.relu = nn.ReLU()
147        self.eps = eps

149    def forward(self, k: torch.Tensor):

151        k = self.dpfp(k)

153        return k / (torch.sum(k, dim=-1, keepdim=True) + self.eps)

155    def dpfp(self, k: torch.Tensor):

160        x = self.relu(torch.cat([k, -k], dim=-1))

163        x_rolled = [x.roll(shifts=i, dims=-1) for i in range(1, self.nu + 1)]

166        x_rolled = torch.cat(x_rolled, dim=-1)

168        x_repeat = torch.cat([x] * self.nu, dim=-1)

174        return x_repeat * x_rolled

177class FastWeightsAttention(Module):

205    def __init__(self, heads: int, d_model: int, dropout_prob: float, phi: DPFP):
206        super().__init__()

209        self.d_k = d_model // heads

211        self.heads = heads

214        self.query = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
215        self.key = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
216        self.value = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)

219        self.interpolation_weight = nn.Sequential(
220            PrepareForMultiHeadAttention(d_model, heads, 1, bias=False),
221            nn.Sigmoid()
222        )

225        self.phi = phi

228        self.output = nn.Linear(d_model, d_model)

230        self.dropout = nn.Dropout(dropout_prob)

232    def forward(self, x: torch.Tensor):

234        seq_len = x.shape[0]

236        query = self.phi(self.query(x))

238        key = self.phi(self.key(x))

240        value = self.value(x)

242        beta = self.interpolation_weight(x)

245        weights = key.new_zeros((key.shape[1], key.shape[2], value.shape[3], key.shape[3]))

247        outputs = []

250        for i in range(seq_len):

252            value_existing = torch.einsum('bhvk,bhk->bhv', weights, key[i])

257            weights = weights + torch.einsum('bhv,bhk->bhvk', beta[i] * (value[i] - value_existing), key[i])

260            y = torch.einsum('bhvk,bhk->bhv', weights, query[i])

263            outputs.append(y.reshape(y.shape[0], -1))

266        x = torch.stack(outputs)

269        return self.output(x)

272class FastWeightsAttentionTransformerLayer(Module):

276    def __init__(self, *,
277                 d_model: int,
278                 attn: FastWeightsAttention,
279                 feed_forward: FeedForward,
280                 dropout_prob: float):
281        super().__init__()

283        self.size = d_model

285        self.attn = attn

287        self.feed_forward = feed_forward

289        self.dropout = nn.Dropout(dropout_prob)

292        self.norm_self_attn = nn.LayerNorm([d_model])
293        self.norm_ff = nn.LayerNorm([d_model])

295    def forward(self, x: torch.Tensor):

297        attn = self.attn(x)

299        x = x + self.dropout(attn)

302        z = self.norm_ff(x)

304        ff = self.feed_forward(z)

306        x = x + self.dropout(ff)

309        return x

312class FastWeightsAttentionTransformer(Module):

316    def __init__(self, layer: FastWeightsAttentionTransformerLayer, n_layers: int):
317        super().__init__()

319        self.layers = clone_module_list(layer, n_layers)

321        self.norm = nn.LayerNorm([layer.size])

323    def forward(self, x: torch.Tensor):
324        for i, layer in enumerate(self.layers):

326            x = layer(x)

329        return self.norm(x)