(self,
dim,
mlp_ratio,
num_heads,
post_norm=False,
causal=False,
activation='quick_gelu',
attn_dropout=0.0,
proj_dropout=0.0,
norm_eps=1e-5)
| 112 | class AttentionBlock(nn.Module): |
| 113 | |
| 114 | def __init__(self, |
| 115 | dim, |
| 116 | mlp_ratio, |
| 117 | num_heads, |
| 118 | post_norm=False, |
| 119 | causal=False, |
| 120 | activation='quick_gelu', |
| 121 | attn_dropout=0.0, |
| 122 | proj_dropout=0.0, |
| 123 | norm_eps=1e-5): |
| 124 | assert activation in ['quick_gelu', 'gelu', 'swi_glu'] |
| 125 | super().__init__() |
| 126 | self.dim = dim |
| 127 | self.mlp_ratio = mlp_ratio |
| 128 | self.num_heads = num_heads |
| 129 | self.post_norm = post_norm |
| 130 | self.causal = causal |
| 131 | self.norm_eps = norm_eps |
| 132 | |
| 133 | # layers |
| 134 | self.norm1 = LayerNorm(dim, eps=norm_eps) |
| 135 | self.attn = SelfAttention(dim, num_heads, causal, attn_dropout, |
| 136 | proj_dropout) |
| 137 | self.norm2 = LayerNorm(dim, eps=norm_eps) |
| 138 | if activation == 'swi_glu': |
| 139 | self.mlp = SwiGLU(dim, int(dim * mlp_ratio)) |
| 140 | else: |
| 141 | self.mlp = nn.Sequential( |
| 142 | nn.Linear(dim, int(dim * mlp_ratio)), |
| 143 | QuickGELU() if activation == 'quick_gelu' else nn.GELU(), |
| 144 | nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) |
| 145 | |
| 146 | def forward(self, x): |
| 147 | if self.post_norm: |
nothing calls this directly
no test coverage detected