(
self,
dim,
num_heads,
mlp_ratio=4.0,
qkv_bias=False,
proj_bias: bool = False,
ffn_bias: bool = True,
drop: float = 0.0,
attn_drop: float = 0.0,
init_values=None,
drop_path: float = 0.0,
act_layer: Callable[..., nn.Module] = nn.GELU,
norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
self_attn_class: Callable[..., nn.Module] = Attention,
cross_attn_class: Callable[..., nn.Module] = CrossAttention,
ffn_layer: Callable[..., nn.Module] = Mlp,
)
| 41 | |
| 42 | class DecoderBlock(nn.Module): |
| 43 | def __init__( |
| 44 | self, |
| 45 | dim, |
| 46 | num_heads, |
| 47 | mlp_ratio=4.0, |
| 48 | qkv_bias=False, |
| 49 | proj_bias: bool = False, |
| 50 | ffn_bias: bool = True, |
| 51 | drop: float = 0.0, |
| 52 | attn_drop: float = 0.0, |
| 53 | init_values=None, |
| 54 | drop_path: float = 0.0, |
| 55 | act_layer: Callable[..., nn.Module] = nn.GELU, |
| 56 | norm_layer: Callable[..., nn.Module] = nn.LayerNorm, |
| 57 | self_attn_class: Callable[..., nn.Module] = Attention, |
| 58 | cross_attn_class: Callable[..., nn.Module] = CrossAttention, |
| 59 | ffn_layer: Callable[..., nn.Module] = Mlp, |
| 60 | ): |
| 61 | super().__init__() |
| 62 | self.norm1 = norm_layer(dim) |
| 63 | self.self_attn = self_attn_class( |
| 64 | dim, |
| 65 | num_heads=num_heads, |
| 66 | qkv_bias=qkv_bias, |
| 67 | proj_bias=proj_bias, |
| 68 | attn_drop=attn_drop, |
| 69 | proj_drop=drop, |
| 70 | ) |
| 71 | self.ls1 = ( |
| 72 | LayerScale(dim, init_values=init_values) if init_values else nn.Identity() |
| 73 | ) |
| 74 | self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 75 | |
| 76 | self.q_norm2 = norm_layer(dim) |
| 77 | self.kv_norm2 = norm_layer(dim) |
| 78 | self.cross_attn = cross_attn_class( |
| 79 | dim, |
| 80 | num_heads=num_heads, |
| 81 | qkv_bias=qkv_bias, |
| 82 | proj_bias=proj_bias, |
| 83 | attn_drop=attn_drop, |
| 84 | proj_drop=drop, |
| 85 | ) |
| 86 | self.ls2 = ( |
| 87 | LayerScale(dim, init_values=init_values) if init_values else nn.Identity() |
| 88 | ) |
| 89 | self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 90 | |
| 91 | self.norm3 = norm_layer(dim) |
| 92 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 93 | self.mlp = ffn_layer( |
| 94 | in_features=dim, |
| 95 | hidden_features=mlp_hidden_dim, |
| 96 | act_layer=act_layer, |
| 97 | drop=drop, |
| 98 | bias=ffn_bias, |
| 99 | ) |
| 100 | self.ls3 = ( |
nothing calls this directly
no test coverage detected