towhee.models.layers.pool_attention

Classes

AttentionPool

A MLP block that contains two linear layers with a normalization layer. the MLP block is used in a transformer model after the attention block. :: Input ↓ Reshape ↓ Pool ↓ Reshape ↓ norm :param thw_shape: the shape of the input tensor (before flattening). :type thw_shape: List :param pool: Pool operation that is applied to the input tensor. If pool is None, return the input tensor. :type pool: Callable :param has_cls_embed: whether the input tensor contains cls token. Pool operation excludes cls token. :type has_cls_embed: bool :param norm: Optional normalization operation applied to tensor after pool. :type norm: Callable.