File size: 1,407 Bytes
961cf0c | 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 | import torch
import torch.nn as nn
class GatedAlteredResidualBlock(nn.Module):
def __init__(self, dim, dropout=0.1):
super(GatedAlteredResidualBlock, self).__init__()
self.norm = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=8, batch_first=False)
self.dropout = nn.Dropout(dropout)
self.gate_linear = nn.Linear(dim, dim)
def forward(self, q, k, v):
attn_output, _ = self.attn(q, k, v)
attn_output = self.dropout(attn_output)
gate = torch.sigmoid(self.gate_linear(q))
out = gate * attn_output + (1 - gate) * q
return self.norm(out)
class StackedGatedResidualBlock(nn.Module):
def __init__(self, dim, depth=4, dropout=0.1):
super(StackedGatedResidualBlock, self).__init__()
self.blocks = nn.ModuleList([
GatedAlteredResidualBlock(dim=dim, dropout=dropout)
for _ in range(depth)
])
def forward(self, q, k=None, v=None):
"""
Permite q = k = v, mas também aceita entrada personalizada.
Args:
q (Tensor): Query, formato (seq_len, batch, dim)
k (Tensor): Key (opcional, default=q)
v (Tensor): Value (opcional, default=q)
"""
if k is None: k = q
if v is None: v = q
for block in self.blocks:
q = block(q, k, v)
return q
|