සම්පීඩ්යතාට්රාන්ස්ෆෝමර්

54from typing import Optional, List
55
56import torch
57import torch.nn.functional as F
58from torch import nn
59
60from labml_helpers.module import Module, TypedModuleList
61from labml_nn.transformers.feed_forward import FeedForward
62from labml_nn.transformers.mha import PrepareForMultiHeadAttention
63from labml_nn.transformers.xl.relative_mha import RelativeMultiHeadAttention
64from labml_nn.utils import clone_module_list

67class Conv1dCompression(Module):

75    def __init__(self, compression_rate: int, d_model: int):

80        super().__init__()
81        self.conv = nn.Conv1d(d_model, d_model, kernel_size=compression_rate, stride=compression_rate)

83    def forward(self, mem: torch.Tensor):

90        mem = mem.permute(1, 2, 0)

92        c_mem = self.conv(mem)

94        return c_mem.permute(2, 0, 1)

97class CompressiveTransformerLayer(Module):

103    def __init__(self, *,
104                 d_model: int,
105                 self_attn: RelativeMultiHeadAttention,
106                 feed_forward: FeedForward,
107                 dropout_prob: float,
108                 compress: Conv1dCompression):

116        super().__init__()
117        self.compress = compress
118        self.size = d_model
119        self.self_attn = self_attn
120        self.feed_forward = feed_forward
121        self.dropout = nn.Dropout(dropout_prob)
122        self.norm_self_attn = nn.LayerNorm([d_model])
123        self.norm_ff = nn.LayerNorm([d_model])

125    def concat_memory(self, z: torch.Tensor, mem: Optional[torch.Tensor], c_mem: Optional[torch.Tensor]):

134        if mem is None:
135            return z

138        if c_mem is not None:
139            mem = torch.cat((c_mem, mem), dim=0)

142        mem = self.norm_self_attn(mem)

144        return torch.cat((mem, z), dim=0)

146    def forward(self, *,
147                x: torch.Tensor,
148                mem: Optional[torch.Tensor],
149                c_mem: Optional[torch.Tensor],
150                mask: torch.Tensor):

160        z = self.norm_self_attn(x)

162        m_z = self.concat_memory(z, mem, c_mem)

164        self_attn = self.self_attn(query=z, key=m_z, value=m_z, mask=mask)

166        x = x + self.dropout(self_attn)

169        z = self.norm_ff(x)

171        ff = self.feed_forward(z)

173        x = x + self.dropout(ff)

176        return x

179class CompressiveTransformer(Module):

186    def __init__(self, layer: CompressiveTransformerLayer, n_layers: int):
187        super().__init__()

189        self.layers = clone_module_list(layer, n_layers)

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

193    def forward(self, x: torch.Tensor, mem: List[torch.Tensor], c_mem: List[torch.Tensor], mask: torch.Tensor):

204        new_mem = []

206        for i, layer in enumerate(self.layers):

208            new_mem.append(x.detach())

210            m = mem[i] if mem else None

212            cm = c_mem[i] if c_mem else None

214            x = layer(x=x, mem=m, c_mem=cm, mask=mask)

216        return self.norm(x), new_mem

219class AttentionReconstructionLoss:

237    def __init__(self, layers: TypedModuleList[CompressiveTransformerLayer]):

241        self.layers = layers
242        self.loss_func = nn.MSELoss()

244    def prepare_for_attn(self, pmha: PrepareForMultiHeadAttention, x: torch.Tensor):

254        head_shape = x.shape[:-1]

257        weight = pmha.linear.weight.detach()
258        bias = pmha.linear.bias.detach() if pmha.linear.bias is not None else None

260        x = F.linear(x, weight, bias)

263        x = x.view(*head_shape, pmha.heads, pmha.d_k)

266        return x

268    def attn(self, layer: RelativeMultiHeadAttention, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):

275        query = self.prepare_for_attn(layer.query, query)
276        key = self.prepare_for_attn(layer.key, key)
277        value = self.prepare_for_attn(layer.value, value)

281        scores = torch.einsum('ibhd,jbhd->ijbh', query, key)

284        scores *= layer.scale

288        attn = layer.softmax(scores)

292        return torch.einsum("ijbh,jbhd->ibhd", attn, value)

294    def norm(self, ln: nn.LayerNorm, x: torch.Tensor):

300        weight = ln.weight.detach() if ln.weight is not None else None
301        bias = ln.bias.detach() if ln.bias is not None else None

304        return F.layer_norm(x, ln.normalized_shape, weight, bias, ln.eps)

306    def calc_loss(self, layer: CompressiveTransformerLayer, h: torch.Tensor, mem: torch.Tensor):

312        h = h.detach()
313        mem = mem.detach()

317        c_mem = layer.compress(mem)

320        h = self.norm(layer.norm_self_attn, h)
321        mem = self.norm(layer.norm_self_attn, mem)
322        c_mem = self.norm(layer.norm_self_attn, c_mem)

325        attn_mem = self.attn(layer.self_attn, h, mem, mem)

327        attn_cmem = self.attn(layer.self_attn, h, c_mem, c_mem)

330        return self.loss_func(attn_cmem, attn_mem)

332    def __call__(self, h: List[torch.Tensor], mem: List[torch.Tensor]):

334        losses = [self.calc_loss(layer, h[n], mem[n]) for n, layer in enumerate(self.layers)]

336        return sum(losses)