| 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) |