(
self,
d_model: int,
n_head: int,
mlp_ratio: float = 4.0,
act_layer: Callable = nn.GELU,
norm_layer: Callable = nn.LayerNorm,
is_cross_attention: bool = False,
)
| 246 | |
| 247 | class VisualAttentionBlock(nn.Module): |
| 248 | def __init__( |
| 249 | self, |
| 250 | d_model: int, |
| 251 | n_head: int, |
| 252 | mlp_ratio: float = 4.0, |
| 253 | act_layer: Callable = nn.GELU, |
| 254 | norm_layer: Callable = nn.LayerNorm, |
| 255 | is_cross_attention: bool = False, |
| 256 | ): |
| 257 | super().__init__() |
| 258 | |
| 259 | self.ln_1 = norm_layer(d_model) |
| 260 | if is_cross_attention: |
| 261 | self.ln_1_kv = norm_layer(d_model) |
| 262 | |
| 263 | self.ln_2 = norm_layer(d_model) |
| 264 | mlp_width = int(d_model * mlp_ratio) |
| 265 | self.attn = VisualAttention(d_model, n_head) |
| 266 | self.mlp = nn.Sequential(OrderedDict([ |
| 267 | ("c_fc", nn.Linear(d_model, mlp_width)), |
| 268 | ("gelu", act_layer()), |
| 269 | ("c_proj", nn.Linear(mlp_width, d_model)) |
| 270 | ])) |
| 271 | |
| 272 | def attention( |
| 273 | self, |
nothing calls this directly
no test coverage detected