diff --git a/timm/models/maxxvit.py b/timm/models/maxxvit.py index 5a164e88..1f8f5f24 100644 --- a/timm/models/maxxvit.py +++ b/timm/models/maxxvit.py @@ -42,6 +42,7 @@ from typing import Callable, Optional, Union, Tuple, List import torch from torch import nn +from torch.jit import Final from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD from timm.layers import Mlp, ConvMlp, DropPath, LayerNorm, ClassifierHead, NormMlpClassifierHead @@ -140,6 +141,8 @@ class MaxxVitCfg: class Attention2d(nn.Module): + fast_attn: Final[bool] + """ multi-head attention for 2D NCHW tensors""" def __init__( self, @@ -208,6 +211,8 @@ class Attention2d(nn.Module): class AttentionCl(nn.Module): """ Channels-last multi-head attention (B, ..., C) """ + fast_attn: Final[bool] + def __init__( self, dim: int, diff --git a/timm/models/vision_transformer.py b/timm/models/vision_transformer.py index 95d87126..dcd0a6fa 100644 --- a/timm/models/vision_transformer.py +++ b/timm/models/vision_transformer.py @@ -33,6 +33,7 @@ import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint +from torch.jit import Final from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD, \ OPENAI_CLIP_MEAN, OPENAI_CLIP_STD @@ -51,6 +52,8 @@ _logger = logging.getLogger(__name__) class Attention(nn.Module): + fast_attn: Final[bool] + def __init__( self, dim, diff --git a/timm/models/vision_transformer_relpos.py b/timm/models/vision_transformer_relpos.py index f9fede53..835b213a 100644 --- a/timm/models/vision_transformer_relpos.py +++ b/timm/models/vision_transformer_relpos.py @@ -11,6 +11,7 @@ from typing import Optional, Tuple import torch import torch.nn as nn +from torch.jit import Final from torch.utils.checkpoint import checkpoint from timm.data import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD @@ -25,6 +26,8 @@ _logger = logging.getLogger(__name__) class RelPosAttention(nn.Module): + fast_attn: Final[bool] + def __init__( self, dim,