| 176 | |
| 177 | |
| 178 | class T5CrossAttention(nn.Module): |
| 179 | |
| 180 | def __init__(self, |
| 181 | dim, |
| 182 | dim_attn, |
| 183 | dim_ffn, |
| 184 | num_heads, |
| 185 | num_buckets, |
| 186 | shared_pos=True, |
| 187 | dropout=0.1): |
| 188 | super(T5CrossAttention, self).__init__() |
| 189 | self.dim = dim |
| 190 | self.dim_attn = dim_attn |
| 191 | self.dim_ffn = dim_ffn |
| 192 | self.num_heads = num_heads |
| 193 | self.num_buckets = num_buckets |
| 194 | self.shared_pos = shared_pos |
| 195 | |
| 196 | # layers |
| 197 | self.norm1 = T5LayerNorm(dim) |
| 198 | self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout) |
| 199 | self.norm2 = T5LayerNorm(dim) |
| 200 | self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout) |
| 201 | self.norm3 = T5LayerNorm(dim) |
| 202 | self.ffn = T5FeedForward(dim, dim_ffn, dropout) |
| 203 | self.pos_embedding = None if shared_pos else T5RelativeEmbedding( |
| 204 | num_buckets, num_heads, bidirectional=False) |
| 205 | |
| 206 | def forward(self, |
| 207 | x, |
| 208 | mask=None, |
| 209 | encoder_states=None, |
| 210 | encoder_mask=None, |
| 211 | pos_bias=None): |
| 212 | e = pos_bias if self.shared_pos else self.pos_embedding( |
| 213 | x.size(1), x.size(1)) |
| 214 | x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e)) |
| 215 | x = fp16_clamp(x + self.cross_attn( |
| 216 | self.norm2(x), context=encoder_states, mask=encoder_mask)) |
| 217 | x = fp16_clamp(x + self.ffn(self.norm3(x))) |
| 218 | return x |
| 219 | |
| 220 | |
| 221 | class T5RelativeEmbedding(nn.Module): |