| 129 | |
| 130 | |
| 131 | class Extractor(nn.Module): |
| 132 | def __init__( |
| 133 | self, d_model, n_head, attn_mask=None, |
| 134 | mlp_factor=4.0, dropout=0.0, drop_path=0.0, |
| 135 | ): |
| 136 | super().__init__() |
| 137 | |
| 138 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 139 | logger.info(f'Drop path rate: {drop_path}') |
| 140 | self.attn = nn.MultiheadAttention(d_model, n_head) |
| 141 | self.ln_1 = nn.LayerNorm(d_model) |
| 142 | d_mlp = round(mlp_factor * d_model) |
| 143 | self.mlp = nn.Sequential(OrderedDict([ |
| 144 | ("c_fc", nn.Linear(d_model, d_mlp)), |
| 145 | ("gelu", QuickGELU()), |
| 146 | ("dropout", nn.Dropout(dropout)), |
| 147 | ("c_proj", nn.Linear(d_mlp, d_model)) |
| 148 | ])) |
| 149 | self.ln_2 = nn.LayerNorm(d_model) |
| 150 | self.ln_3 = nn.LayerNorm(d_model) |
| 151 | self.attn_mask = attn_mask |
| 152 | |
| 153 | # zero init |
| 154 | nn.init.xavier_uniform_(self.attn.in_proj_weight) |
| 155 | nn.init.constant_(self.attn.out_proj.weight, 0.) |
| 156 | nn.init.constant_(self.attn.out_proj.bias, 0.) |
| 157 | nn.init.xavier_uniform_(self.mlp[0].weight) |
| 158 | nn.init.constant_(self.mlp[-1].weight, 0.) |
| 159 | nn.init.constant_(self.mlp[-1].bias, 0.) |
| 160 | |
| 161 | def attention(self, x, y): |
| 162 | d_model = self.ln_1.weight.size(0) |
| 163 | q = (x @ self.attn.in_proj_weight[:d_model].T) + self.attn.in_proj_bias[:d_model] |
| 164 | |
| 165 | k = (y @ self.attn.in_proj_weight[d_model:-d_model].T) + self.attn.in_proj_bias[d_model:-d_model] |
| 166 | v = (y @ self.attn.in_proj_weight[-d_model:].T) + self.attn.in_proj_bias[-d_model:] |
| 167 | Tx, Ty, N = q.size(0), k.size(0), q.size(1) |
| 168 | q = q.view(Tx, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3) |
| 169 | k = k.view(Ty, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3) |
| 170 | v = v.view(Ty, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3) |
| 171 | aff = (q @ k.transpose(-2, -1) / (self.attn.head_dim ** 0.5)) |
| 172 | |
| 173 | aff = aff.softmax(dim=-1) |
| 174 | out = aff @ v |
| 175 | out = out.permute(2, 0, 1, 3).flatten(2) |
| 176 | out = self.attn.out_proj(out) |
| 177 | return out |
| 178 | |
| 179 | def forward(self, x, y): |
| 180 | x = x + self.drop_path(self.attention(self.ln_1(x), self.ln_3(y))) |
| 181 | x = x + self.drop_path(self.mlp(self.ln_2(x))) |
| 182 | return x |
| 183 | |
| 184 | |
| 185 | class Transformer(nn.Module): |