| 165 | |
| 166 | |
| 167 | class ResidualAttentionBlock(nn.Module): |
| 168 | def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None): |
| 169 | super().__init__() |
| 170 | |
| 171 | self.attn = MultiheadAttention(d_model, n_head) |
| 172 | self.ln_1 = LayerNorm(d_model) |
| 173 | self.mlp = nn.Sequential(OrderedDict([ |
| 174 | ("c_fc", nn.Linear(d_model, d_model * 4)), |
| 175 | ("gelu", QuickGELU()), |
| 176 | ("c_proj", nn.Linear(d_model * 4, d_model)) |
| 177 | ])) |
| 178 | self.ln_2 = LayerNorm(d_model) |
| 179 | self.attn_mask = attn_mask |
| 180 | |
| 181 | self.attn_probs = None |
| 182 | self.attn_grad = None |
| 183 | |
| 184 | def set_attn_probs(self, attn_probs): |
| 185 | self.attn_probs = attn_probs |
| 186 | |
| 187 | def set_attn_grad(self, attn_grad): |
| 188 | self.attn_grad = attn_grad |
| 189 | |
| 190 | def attention(self, x: torch.Tensor): |
| 191 | self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None |
| 192 | return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask, attention_probs_forward_hook=self.set_attn_probs, |
| 193 | attention_probs_backwards_hook=self.set_attn_grad)[0] |
| 194 | |
| 195 | def forward(self, x: torch.Tensor): |
| 196 | x = x + self.attention(self.ln_1(x)) |
| 197 | x = x + self.mlp(self.ln_2(x)) |
| 198 | return x |
| 199 | |
| 200 | |
| 201 | class Transformer(nn.Module): |