Merge pull request #250 from rwightman/vision_transformer

Vision Transformer

Author: rwightman

Diff for: README.md

@@ -2,6 +2,12 @@
 
 ## What's New
 
+### Oct 13, 2020
+* Initial impl of Vision Transformer models. Both patch and hybrid (CNN backbone) variants. Currently trying to train...
+* Adafactor and AdaHessian (FP32 only, no AMP) optimizers
+* EdgeTPU-M (`efficientnet_em`) model trained in PyTorch, 79.3 top-1
+* Pip release, doc updates pending a few more changes...
+
 ### Sept 18, 2020
 * New ResNet 'D' weights. 72.7 (top-1) ResNet-18-D, 77.1 ResNet-34-D, 80.5 ResNet-50-D
 * Added a few untrained defs for other ResNet models (66D, 101D, 152D, 200/200D)
@@ -124,6 +130,7 @@ A full version of the list below with source links can be found in the [document
 * SelecSLS - https://arxiv.org/abs/1907.00837
 * Selective Kernel Networks - https://arxiv.org/abs/1903.06586
 * TResNet - https://arxiv.org/abs/2003.13630
+* Vision Transformer - https://openreview.net/forum?id=YicbFdNTTy
 * VovNet V2 and V1 - https://arxiv.org/abs/1911.06667
 * Xception - https://arxiv.org/abs/1610.02357
 * Xception (Modified Aligned, Gluon) - https://arxiv.org/abs/1802.02611
@@ -162,6 +169,8 @@ Several (less common) features that I often utilize in my projects are included.
     * `lookahead` adapted from impl by [Liam](https://github.com/alphadl/lookahead.pytorch) (https://arxiv.org/abs/1907.08610)
     * `fused<name>` optimizers by name with [NVIDIA Apex](https://github.com/NVIDIA/apex/tree/master/apex/optimizers) installed
     * `adamp` and `sgdp` by [Naver ClovAI](https://github.com/clovaai) (https://arxiv.org/abs/2006.08217)
+    * `adafactor` adapted from [FAIRSeq impl](https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py) (https://arxiv.org/abs/1804.04235)
+    * `adahessian` by [David Samuel](https://github.com/davda54/ada-hessian) (https://arxiv.org/abs/2006.00719)
 * Random Erasing from [Zhun Zhong](https://github.com/zhunzhong07/Random-Erasing/blob/master/transforms.py)  (https://arxiv.org/abs/1708.04896)
 * Mixup (https://arxiv.org/abs/1710.09412)
 * CutMix (https://arxiv.org/abs/1905.04899)

Diff for: avg_checkpoints.py

@@ -103,7 +103,11 @@ def main():
         v = v.clamp(float32_info.min, float32_info.max)
         final_state_dict[k] = v.to(dtype=torch.float32)
 
-    torch.save(final_state_dict, args.output)
+    try:
+        torch.save(final_state_dict, args.output, _use_new_zipfile_serialization=False)
+    except:
+        torch.save(final_state_dict, args.output)
+
     with open(args.output, 'rb') as f:
         sha_hash = hashlib.sha256(f.read()).hexdigest()
     print("=> Saved state_dict to '{}, SHA256: {}'".format(args.output, sha_hash))

Diff for: clean_checkpoint.py

@@ -57,7 +57,11 @@ def main():
             new_state_dict[name] = v
         print("=> Loaded state_dict from '{}'".format(args.checkpoint))
 
-        torch.save(new_state_dict, _TEMP_NAME)
+        try:
+            torch.save(new_state_dict, _TEMP_NAME, _use_new_zipfile_serialization=False)
+        except:
+            torch.save(new_state_dict, _TEMP_NAME)
+
         with open(_TEMP_NAME, 'rb') as f:
             sha_hash = hashlib.sha256(f.read()).hexdigest()
 

Diff for: 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
@@ -68,7 +68,7 @@ def test_model_backward(model_name, batch_size):
 
 
 @pytest.mark.timeout(120)
-@pytest.mark.parametrize('model_name', list_models())
+@pytest.mark.parametrize('model_name', list_models(exclude_filters=['vit_*']))
 @pytest.mark.parametrize('batch_size', [1])
 def test_model_default_cfgs(model_name, batch_size):
     """Run a single forward pass with each model"""

Diff for: 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 *

Diff for: 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

Diff for: 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
 

Diff for: 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

Diff for: 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
 
 
 

Diff for: 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):

Diff for: 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):

Diff for: 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