(self, x, hw_shape, identity=None)
| 147 | self.norm = build_norm_layer(norm_cfg, embed_dims)[1] |
| 148 | |
| 149 | def forward(self, x, hw_shape, identity=None): |
| 150 | |
| 151 | x_q = x |
| 152 | if self.sr_ratio > 1: |
| 153 | x_kv = nlc_to_nchw(x, hw_shape) |
| 154 | x_kv = self.sr(x_kv) |
| 155 | x_kv = nchw_to_nlc(x_kv) |
| 156 | x_kv = self.norm(x_kv) |
| 157 | else: |
| 158 | x_kv = x |
| 159 | |
| 160 | if identity is None: |
| 161 | identity = x_q |
| 162 | |
| 163 | # `need_weights=True` will let nn.MultiHeadAttention |
| 164 | # `return attn_output, attn_output_weights.sum(dim=1) / num_heads` |
| 165 | # The `attn_output_weights.sum(dim=1)` may cause cuda error. So, we set |
| 166 | # `need_weights=False` to ignore `attn_output_weights.sum(dim=1)`. |
| 167 | # This issue - `https://github.com/pytorch/pytorch/issues/37583` report |
| 168 | # the error that large scale tensor sum operation may cause cuda error. |
| 169 | out = self.attn(query=x_q, key=x_kv, value=x_kv, need_weights=False)[0] |
| 170 | |
| 171 | return identity + self.dropout_layer(self.proj_drop(out)) |
| 172 | |
| 173 | |
| 174 | class TransformerEncoderLayer(BaseModule): |
nothing calls this directly
no outgoing calls
no test coverage detected