The transformer encoder layer from CLIP.
| 25 | |
| 26 | |
| 27 | class CLIPEncoderLayer(nn.Module): |
| 28 | """The transformer encoder layer from CLIP.""" |
| 29 | |
| 30 | def __init__(self, model_dims: int, num_heads: int, activation: str): |
| 31 | super().__init__() |
| 32 | |
| 33 | self.layer_norm1 = nn.LayerNorm(model_dims) |
| 34 | self.layer_norm2 = nn.LayerNorm(model_dims) |
| 35 | |
| 36 | self.attention = nn.MultiHeadAttention(model_dims, num_heads) |
| 37 | # Add biases to the attention projections to match CLIP |
| 38 | self.attention.query_proj.bias = mx.zeros(model_dims) |
| 39 | self.attention.key_proj.bias = mx.zeros(model_dims) |
| 40 | self.attention.value_proj.bias = mx.zeros(model_dims) |
| 41 | self.attention.out_proj.bias = mx.zeros(model_dims) |
| 42 | |
| 43 | self.linear1 = nn.Linear(model_dims, 4 * model_dims) |
| 44 | self.linear2 = nn.Linear(4 * model_dims, model_dims) |
| 45 | |
| 46 | self.act = _ACTIVATIONS[activation] |
| 47 | |
| 48 | def __call__(self, x, attn_mask=None): |
| 49 | y = self.layer_norm1(x) |
| 50 | y = self.attention(y, y, y, attn_mask) |
| 51 | x = y + x |
| 52 | |
| 53 | y = self.layer_norm2(x) |
| 54 | y = self.linear1(y) |
| 55 | y = self.act(y) |
| 56 | y = self.linear2(y) |
| 57 | x = y + x |
| 58 | |
| 59 | return x |
| 60 | |
| 61 | |
| 62 | class CLIPTextModel(nn.Module): |