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