| 111 | |
| 112 | |
| 113 | class T5SelfAttention(nn.Module): |
| 114 | |
| 115 | def __init__(self, |
| 116 | dim, |
| 117 | dim_attn, |
| 118 | dim_ffn, |
| 119 | num_heads, |
| 120 | num_buckets, |
| 121 | shared_pos=True, |
| 122 | dropout=0.1): |
| 123 | super(T5SelfAttention, self).__init__() |
| 124 | self.dim = dim |
| 125 | self.dim_attn = dim_attn |
| 126 | self.dim_ffn = dim_ffn |
| 127 | self.num_heads = num_heads |
| 128 | self.num_buckets = num_buckets |
| 129 | self.shared_pos = shared_pos |
| 130 | |
| 131 | # layers |
| 132 | self.norm1 = T5LayerNorm(dim) |
| 133 | self.attn = T5Attention(dim, dim_attn, num_heads, dropout) |
| 134 | self.norm2 = T5LayerNorm(dim) |
| 135 | self.ffn = T5FeedForward(dim, dim_ffn, dropout) |
| 136 | self.pos_embedding = None if shared_pos else T5RelativeEmbedding( |
| 137 | num_buckets, num_heads, bidirectional=True) |
| 138 | |
| 139 | def forward(self, x, mask=None, pos_bias=None): |
| 140 | e = pos_bias if self.shared_pos else self.pos_embedding( |
| 141 | x.size(1), x.size(1)) |
| 142 | x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) |
| 143 | x = fp16_clamp(x + self.ffn(self.norm2(x))) |
| 144 | return x |
| 145 | |
| 146 | |
| 147 | class T5RelativeEmbedding(nn.Module): |