Initial Vision Transformer impl w/ patch and hybrid variants. Refactor tuple helpers.

Author: rwightman

tests/test_models.py

@@ -15,9 +15,9 @@
 
 if 'GITHUB_ACTIONS' in os.environ:  # and 'Linux' in platform.system():
     # GitHub Linux runner is slower and hits memory limits sooner than MacOS, exclude bigger models
-    EXCLUDE_FILTERS = ['*efficientnet_l2*', '*resnext101_32x48d']
+    EXCLUDE_FILTERS = ['*efficientnet_l2*', '*resnext101_32x48d', 'vit_*']
 else:
-    EXCLUDE_FILTERS = []
+    EXCLUDE_FILTERS = ['vit_*']
 MAX_FWD_SIZE = 384
 MAX_BWD_SIZE = 128
 MAX_FWD_FEAT_SIZE = 448

timm/models/__init__.py

@@ -21,6 +21,7 @@
 from .senet import *
 from .sknet import *
 from .tresnet import *
+from .vision_transformer import *
 from .vovnet import *
 from .xception import *
 from .xception_aligned import *

timm/models/layers/__init__.py

@@ -16,6 +16,7 @@
 from .drop import DropBlock2d, DropPath, drop_block_2d, drop_path
 from .eca import EcaModule, CecaModule
 from .evo_norm import EvoNormBatch2d, EvoNormSample2d
+from .helpers import to_ntuple, to_2tuple, to_3tuple, to_4tuple
 from .inplace_abn import InplaceAbn
 from .mixed_conv2d import MixedConv2d
 from .norm_act import BatchNormAct2d

timm/models/layers/cond_conv2d.py

@@ -13,7 +13,7 @@
 from torch import nn as nn
 from torch.nn import functional as F
 
-from .helpers import tup_pair
+from .helpers import to_2tuple
 from .conv2d_same import conv2d_same
 from .padding import get_padding_value
 
@@ -46,13 +46,13 @@ def __init__(self, in_channels, out_channels, kernel_size=3,
 
         self.in_channels = in_channels
         self.out_channels = out_channels
-        self.kernel_size = tup_pair(kernel_size)
-        self.stride = tup_pair(stride)
+        self.kernel_size = to_2tuple(kernel_size)
+        self.stride = to_2tuple(stride)
         padding_val, is_padding_dynamic = get_padding_value(
             padding, kernel_size, stride=stride, dilation=dilation)
         self.dynamic_padding = is_padding_dynamic  # if in forward to work with torchscript
-        self.padding = tup_pair(padding_val)
-        self.dilation = tup_pair(dilation)
+        self.padding = to_2tuple(padding_val)
+        self.dilation = to_2tuple(dilation)
         self.groups = groups
         self.num_experts = num_experts
 

timm/models/layers/drop.py

@@ -150,7 +150,8 @@ def drop_path(x, drop_prob: float = 0., training: bool = False):
     if drop_prob == 0. or not training:
         return x
     keep_prob = 1 - drop_prob
-    random_tensor = keep_prob + torch.rand((x.size()[0], 1, 1, 1), dtype=x.dtype, device=x.device)
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
     random_tensor.floor_()  # binarize
     output = x.div(keep_prob) * random_tensor
     return output

timm/models/layers/helpers.py

@@ -15,11 +15,11 @@ def parse(x):
     return parse
 
 
-tup_single = _ntuple(1)
-tup_pair = _ntuple(2)
-tup_triple = _ntuple(3)
-tup_quadruple = _ntuple(4)
-ntup = _ntuple
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
 
 
 

timm/models/layers/median_pool.py

@@ -3,7 +3,7 @@
 """
 import torch.nn as nn
 import torch.nn.functional as F
-from .helpers import tup_pair, tup_quadruple
+from .helpers import to_2tuple, to_4tuple
 
 
 class MedianPool2d(nn.Module):
@@ -17,9 +17,9 @@ class MedianPool2d(nn.Module):
     """
     def __init__(self, kernel_size=3, stride=1, padding=0, same=False):
         super(MedianPool2d, self).__init__()
-        self.k = tup_pair(kernel_size)
-        self.stride = tup_pair(stride)
-        self.padding = tup_quadruple(padding)  # convert to l, r, t, b
+        self.k = to_2tuple(kernel_size)
+        self.stride = to_2tuple(stride)
+        self.padding = to_4tuple(padding)  # convert to l, r, t, b
         self.same = same
 
     def _padding(self, x):

timm/models/layers/pool2d_same.py

@@ -7,7 +7,7 @@
 import torch.nn.functional as F
 from typing import List, Tuple, Optional
 
-from .helpers import tup_pair
+from .helpers import to_2tuple
 from .padding import pad_same, get_padding_value
 
 
@@ -22,8 +22,8 @@ class AvgPool2dSame(nn.AvgPool2d):
     """ Tensorflow like 'SAME' wrapper for 2D average pooling
     """
     def __init__(self, kernel_size: int, stride=None, padding=0, ceil_mode=False, count_include_pad=True):
-        kernel_size = tup_pair(kernel_size)
-        stride = tup_pair(stride)
+        kernel_size = to_2tuple(kernel_size)
+        stride = to_2tuple(stride)
         super(AvgPool2dSame, self).__init__(kernel_size, stride, (0, 0), ceil_mode, count_include_pad)
 
     def forward(self, x):
@@ -42,9 +42,9 @@ class MaxPool2dSame(nn.MaxPool2d):
     """ Tensorflow like 'SAME' wrapper for 2D max pooling
     """
     def __init__(self, kernel_size: int, stride=None, padding=0, dilation=1, ceil_mode=False, count_include_pad=True):
-        kernel_size = tup_pair(kernel_size)
-        stride = tup_pair(stride)
-        dilation = tup_pair(dilation)
+        kernel_size = to_2tuple(kernel_size)
+        stride = to_2tuple(stride)
+        dilation = to_2tuple(dilation)
         super(MaxPool2dSame, self).__init__(kernel_size, stride, (0, 0), dilation, ceil_mode, count_include_pad)
 
     def forward(self, x):

timm/models/rexnet.py

@@ -17,6 +17,7 @@
 from .helpers import build_model_with_cfg
 from .layers import ClassifierHead, create_act_layer, ConvBnAct, DropPath
 from .registry import register_model
+from .efficientnet_builder import efficientnet_init_weights
 
 
 def _cfg(url=''):
@@ -186,7 +187,7 @@ def __init__(self, in_chans=3, num_classes=1000, global_pool='avg', output_strid
 
         self.head = ClassifierHead(self.num_features, num_classes, global_pool, drop_rate)
 
-        # FIXME weight init, the original appears to use PyTorch defaults
+        efficientnet_init_weights(self)
 
     def get_classifier(self):
         return self.head.fc