VisionLanguageGroup/MicroscopyMatching

1from .mlp import MLP2 3from torch import nn4 5 6class TransformerEncoder(nn.Module):7 8    def __init__(9        self,10        num_layers: int,11        emb_dim: int,12        num_heads: int,13        dropout: float,14        layer_norm_eps: float,15        mlp_factor: int,16        norm_first: bool,17        activation: nn.Module,18        norm: bool,19    ):20 21        super(TransformerEncoder, self).__init__()22 23        self.layers = nn.ModuleList([24            TransformerEncoderLayer(25                emb_dim, num_heads, dropout, layer_norm_eps,26                mlp_factor, norm_first, activation27            ) for _ in range(num_layers)28        ])29 30        self.norm = nn.LayerNorm(emb_dim, layer_norm_eps) if norm else nn.Identity()31 32    def forward(self, src, pos_emb, src_mask, src_key_padding_mask):33        output = src34        for layer in self.layers:35            output = layer(output, pos_emb, src_mask, src_key_padding_mask)36        return self.norm(output)37 38 39class TransformerEncoderLayer(nn.Module):40 41    def __init__(42        self,43        emb_dim: int,44        num_heads: int,45        dropout: float,46        layer_norm_eps: float,47        mlp_factor: int,48        norm_first: bool,49        activation: nn.Module,50    ):51        super(TransformerEncoderLayer, self).__init__()52 53        self.norm_first = norm_first54 55        self.norm1 = nn.LayerNorm(emb_dim, layer_norm_eps)56        self.norm2 = nn.LayerNorm(emb_dim, layer_norm_eps)57        self.dropout1 = nn.Dropout(dropout)58        self.dropout2 = nn.Dropout(dropout)59 60        self.self_attn = nn.MultiheadAttention(61            emb_dim, num_heads, dropout62        )63        self.mlp = MLP(emb_dim, mlp_factor * emb_dim, dropout, activation)64 65    def with_emb(self, x, emb):66        return x if emb is None else x + emb67 68    def forward(self, src, pos_emb, src_mask, src_key_padding_mask):69        if self.norm_first:70            src_norm = self.norm1(src)71            q = k = src_norm + pos_emb72            src = src + self.dropout1(self.self_attn(73                query=q,74                key=k,75                value=src_norm,76                attn_mask=src_mask,77                key_padding_mask=src_key_padding_mask78            )[0])79 80            src_norm = self.norm2(src)81            src = src + self.dropout2(self.mlp(src_norm))82        else:83            q = k = src + pos_emb84            src = self.norm1(src + self.dropout1(self.self_attn(85                query=q,86                key=k,87                value=src,88                attn_mask=src_mask,89                key_padding_mask=src_key_padding_mask90            )[0]))91 92            src = self.norm2(src + self.dropout2(self.mlp(src)))93 94        return src95