vit draft (base,small,tiny), small fix to cifar

Author: jiuntian

configs.py

@@ -339,6 +339,7 @@ def dataset(config, filename, transform_mode,
         if dataset_name in ['cifar10', 'cifar10_II', 'cifar100']:
             d = datasets.cifar(nclass, transform=transform, filename=filename, evaluation_protocol=ep, reset=reset,
                                remove_train_from_db=remove_train_from_db, extra_dataset=extra_dataset)
+            logging.info(f'Number of data: {len(d.data)}')
             logging.info(f'Augmentation for {transform_mode}: {transform.transforms}')
         else:
             raise NotImplementedError(f"Not implementation for {dataset_name}")

models/architectures/helper.py

@@ -4,6 +4,7 @@
 from models.backbone.linear import LinearBackbone
 from models.backbone.resnet import ResNetBackbone
 from models.backbone.vgg import VGGBackbone
+from models.backbone.vit import ViTBackbone
 from models.layers.activation import SignHashLayer, StochasticBinaryLayer
 from models.layers.bihalf import BiHalfLayer
 from models.layers.zm import MeanOnlyBatchNorm
@@ -36,6 +37,15 @@ def get_backbone(backbone, nbit, nclass, pretrained, freeze_weight, **kwargs):
                            vgg_size='vgg16bn', freeze_weight=freeze_weight, **kwargs)
     elif backbone == 'linear':
         return LinearBackbone(nclass=nclass, nbit=nbit, **kwargs)
+    elif backbone == 'vit':
+        return ViTBackbone(nbit=nbit, nclass=nclass, vit_name='vit_base_patch16_224',
+                       pretrained=pretrained, freeze_weight=freeze_weight, **kwargs)
+    elif backbone == 'vittiny':
+        return ViTBackbone(nbit=nbit, nclass=nclass, vit_name='vit_tiny_patch16_224',
+                       pretrained=pretrained, freeze_weight=freeze_weight, **kwargs)
+    elif backbone == 'vitsmall':
+        return ViTBackbone(nbit=nbit, nclass=nclass, vit_name='vit_small_patch16_224',
+                       pretrained=pretrained, freeze_weight=freeze_weight, **kwargs)
     else:
         raise NotImplementedError('The backbone not implemented.')
 

models/backbone/vit.py

@@ -0,0 +1,49 @@
+from abc import ABC
+from typing import List
+
+import timm
+import torch
+from torch import nn
+
+from models.backbone.base_backbone import BaseBackbone
+
+
+class ViTBackbone(BaseBackbone):
+    def __init__(self, nbit, nclass, vit_name, pretrained=False, freeze_weight=False, **kwargs):
+        super(ViTBackbone, self).__init__()
+
+        model = timm.create_model(vit_name, pretrained=pretrained)
+
+        self.patch_embed = model.patch_embed
+        self.cls_token = model.cls_token
+        self.pos_embed = model.pos_embed
+        self.pos_drop = model.pos_drop
+        self.blocks = model.blocks
+        self.norm = model.norm
+        self.pre_logits = model.pre_logits
+        self.head = model.head  # no need train as features_params because not using
+
+        self.in_features = model.head.in_features
+        self.nbit = nbit
+        self.nclass = nclass
+
+        assert freeze_weight is False, \
+            'freeze_weight in backbone deprecated. Use --backbone-lr-scale=0 to freeze backbone'
+
+    def get_features_params(self) -> List:
+        return list(self.parameters())
+
+    def get_hash_params(self) -> List:
+        raise NotImplementedError('no hash layer in backbone')
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+
+        cls_token = self.cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_token, x), dim=1)
+
+        x = self.pos_drop(x + self.pos_embed)
+        x = self.blocks(x)
+        x = self.norm(x)
+
+        return self.pre_logits(x[:, 0])

requirements.txt

@@ -1,5 +1,6 @@
 torch~=1.10.0
 torchvision~=0.11.0
+timm~=0.5.4
 tqdm
 opencv-python
 scikit-learn

utils/datasets.py

@@ -616,7 +616,9 @@ def cifar(nclass, **kwargs):
                    transform=transform, target_transform=one_hot(int(nclass)),
                    train=True, download=True)
     traind = IndexDatasetWrapper(traind)
-    testd = CIFAR(f'data/cifar{nclass}', train=False, download=True)
+    testd = CIFAR(f'data/cifar{nclass}',
+                  transform=transform, target_transform=one_hot(int(nclass)),
+                  train=False, download=True)
     testd = IndexDatasetWrapper(testd)
 
     if ep == 2:  # using orig train and test