[20230608 v0.5.1] Fix compatibility issues

qtcls/models/cait.py

@@ -7,7 +7,7 @@
 
 import torch
 import torch.nn as nn
-from timm.models.layers import PatchEmbed, Mlp, DropPath, trunc_normal_
+from timm.layers import PatchEmbed, Mlp, DropPath, trunc_normal_
 
 __all__ = [
     'Cait',

qtcls/models/deit.py

@@ -5,8 +5,10 @@
 # -------------------------------------------------------------------------------
 
 import torch
-from timm.models.vision_transformer import VisionTransformer, trunc_normal_
-from torch import nn as nn
+import torch.nn as nn
+from timm.layers import trunc_normal_
+
+from .vision_transformer_timm import VisionTransformer
 
 __all__ = [
     'VisionTransformerDistilled',

qtcls/models/levit.py

@@ -8,8 +8,7 @@
 
 import torch
 import torch.nn as nn
-from timm.models.layers import to_ntuple, get_act_layer
-from timm.models.vision_transformer import trunc_normal_
+from timm.layers import to_ntuple, get_act_layer, trunc_normal_
 
 __all__ = [
     'Levit',

qtcls/models/mlp_mixer.py

@@ -8,8 +8,8 @@
 
 import torch
 import torch.nn as nn
-from timm.models.helpers import named_apply
-from timm.models.layers import PatchEmbed, Mlp, GluMlp, GatedMlp, DropPath, lecun_normal_, to_2tuple
+from timm.layers import PatchEmbed, Mlp, GluMlp, GatedMlp, DropPath, lecun_normal_, to_2tuple
+from timm.models import named_apply
 
 __all__ = [
     'MlpMixer',

qtcls/models/pvt.py

@@ -7,7 +7,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.layers import DropPath, to_2tuple, trunc_normal_
 
 __all__ = [
     'PyramidVisionTransformer',

qtcls/models/swin_transformer.py

@@ -8,7 +8,7 @@
 
 import torch
 import torch.nn as nn
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.layers import DropPath, to_2tuple, trunc_normal_
 
 __all__ = [
     'SwinTransformer',

qtcls/models/swin_transformer_v2.py

@@ -10,8 +10,8 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from timm.models.fx_features import register_notrace_function
-from timm.models.layers import PatchEmbed, Mlp, DropPath, to_2tuple, trunc_normal_, _assert
+from timm.layers import PatchEmbed, Mlp, DropPath, to_2tuple, trunc_normal_, _assert
+from timm.models import register_notrace_function
 
 __all__ = [
     'SwinTransformerV2',

qtcls/models/tnt.py

@@ -6,9 +6,7 @@
 
 import torch
 import torch.nn as nn
-from timm.models.layers import DropPath, Mlp, trunc_normal_
-from timm.models.layers import _assert
-from timm.models.layers.helpers import to_2tuple
+from timm.layers import _assert, to_2tuple, DropPath, Mlp, trunc_normal_
 
 __all__ = [
     'TNT',

qtcls/models/twins.py

@@ -9,8 +9,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from timm.models.layers import DropPath, Mlp, to_2tuple, trunc_normal_
-from timm.models.vision_transformer import Attention
+from timm.layers import DropPath, Mlp, to_2tuple, trunc_normal_
 
 __all__ = [
     'Twins',
@@ -25,6 +24,34 @@
 Size_ = Tuple[int, int]
 
 
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
 class LocallyGroupedAttn(nn.Module):
     """ LSA: self attention within a group
     """

qtcls/models/vision_transformer_timm.py

@@ -8,10 +8,8 @@
 
 import torch
 import torch.nn as nn
-from timm.models.helpers import named_apply
-from timm.models.layers import trunc_normal_, lecun_normal_
-from timm.models.layers.helpers import to_2tuple
-from timm.models.layers.trace_utils import _assert
+from timm.layers import _assert, to_2tuple, trunc_normal_, lecun_normal_
+from timm.models import named_apply
 
 __all__ = [
     'VisionTransformer',

requirements.txt

@@ -5,4 +5,4 @@ Pillow
 scikit-learn
 scipy
 termcolor
-timm
+timm==0.9.2