Implements one encoder layer in Vision Transformer. Args: embed_dims (int): The feature dimension. num_heads (int): Parallel attention heads. feedforward_channels (int): The hidden dimension for FFNs. drop_rate (float): Probability of an element to be zeroed
| 18 | |
| 19 | |
| 20 | class TransformerEncoderLayer(BaseModule): |
| 21 | """Implements one encoder layer in Vision Transformer. |
| 22 | |
| 23 | Args: |
| 24 | embed_dims (int): The feature dimension. |
| 25 | num_heads (int): Parallel attention heads. |
| 26 | feedforward_channels (int): The hidden dimension for FFNs. |
| 27 | drop_rate (float): Probability of an element to be zeroed |
| 28 | after the feed forward layer. Default: 0.0. |
| 29 | attn_drop_rate (float): The drop out rate for attention layer. |
| 30 | Default: 0.0. |
| 31 | drop_path_rate (float): stochastic depth rate. Default 0.0. |
| 32 | num_fcs (int): The number of fully-connected layers for FFNs. |
| 33 | Default: 2. |
| 34 | qkv_bias (bool): enable bias for qkv if True. Default: True |
| 35 | act_cfg (dict): The activation config for FFNs. |
| 36 | Defalut: dict(type='GELU'). |
| 37 | norm_cfg (dict): Config dict for normalization layer. |
| 38 | Default: dict(type='LN'). |
| 39 | batch_first (bool): Key, Query and Value are shape of |
| 40 | (batch, n, embed_dim) |
| 41 | or (n, batch, embed_dim). Default: True. |
| 42 | """ |
| 43 | |
| 44 | def __init__(self, |
| 45 | embed_dims, |
| 46 | num_heads, |
| 47 | feedforward_channels, |
| 48 | drop_rate=0., |
| 49 | attn_drop_rate=0., |
| 50 | drop_path_rate=0., |
| 51 | num_fcs=2, |
| 52 | qkv_bias=True, |
| 53 | act_cfg=dict(type='GELU'), |
| 54 | norm_cfg=dict(type='LN'), |
| 55 | batch_first=True): |
| 56 | super(TransformerEncoderLayer, self).__init__() |
| 57 | |
| 58 | self.norm1_name, norm1 = build_norm_layer( |
| 59 | norm_cfg, embed_dims, postfix=1) |
| 60 | self.add_module(self.norm1_name, norm1) |
| 61 | |
| 62 | self.attn = MultiheadAttention( |
| 63 | embed_dims=embed_dims, |
| 64 | num_heads=num_heads, |
| 65 | attn_drop=attn_drop_rate, |
| 66 | proj_drop=drop_rate, |
| 67 | dropout_layer=dict(type='DropPath', drop_prob=drop_path_rate), |
| 68 | batch_first=batch_first, |
| 69 | bias=qkv_bias) |
| 70 | |
| 71 | self.norm2_name, norm2 = build_norm_layer( |
| 72 | norm_cfg, embed_dims, postfix=2) |
| 73 | self.add_module(self.norm2_name, norm2) |
| 74 | |
| 75 | self.ffn = FFN( |
| 76 | embed_dims=embed_dims, |
| 77 | feedforward_channels=feedforward_channels, |