update fine-tuning instructions

zihangJiang · Jun 16, 2021 · aa438ef · aa438ef
1 parent 73cf498
commit aa438ef
Show file tree

Hide file tree

Showing 20 changed files with 20 additions and 15 deletions.
diff --git a/README.md b/README.md
@@ -10,7 +10,7 @@ Our codes are based on the [pytorch-image-models](https://github.com/rwightman/p
 
 ### Update
 
-**2021.6: Support `pip install tlmm` to use our Token Labeling for image models.**
+**2021.6: Support `pip install tlt` to use our Token Labeling Toolbox for image models.**
 
 **2021.6: Release training code and segmentation model.**
 
@@ -99,6 +99,11 @@ To Fine-tune the pre-trained LV-ViT-S on images with 384x384 resolution:
 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 /path/to/imagenet --model lvvit_s -b 64 --apex-amp --img-size 384 --drop-path 0.1 --token-label --token-label-data /path/to/label_data --token-label-size 24 --lr 5.e-6 --min-lr 5.e-6 --weight-decay 1.e-8 --finetune /path/to/checkpoint
 ```
 
+To Fine-tune the pre-trained LV-ViT-S on other datasets without token labeling:
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 /path/to/dataset --model lvvit_s -b 64 --apex-amp --img-size 224 --drop-path 0.1 --token-label --token-label-size 14 --dense-weight 0.0 --num-classes $NUM_CLASSES --finetune /path/to/checkpoint
+```
+
 ### Segmentation
 
 Our Segmentation model are fully based upon the [MMSegmentation](https://github.com/open-mmlab/mmsegmentation) Toolkit. The model and config files are under `seg/` folder which follow the same folder structure. You can simply drop in these file to get start.

diff --git a/flops_computation.py b/flops_computation.py
@@ -1,9 +1,9 @@
-import tlmm.models
+import tlt.models
 # summary of model flops and parameters
 
-model_list = [tlmm.models.lvvit_s,
-              tlmm.models.lvvit_m,
-              tlmm.models.lvvit_l]
+model_list = [tlt.models.lvvit_s,
+              tlt.models.lvvit_m,
+              tlt.models.lvvit_l]
 
 img_size_list=[224,288,384,448]
 

diff --git a/main.py b/main.py
@@ -24,10 +24,10 @@
 from timm.scheduler import create_scheduler
 from timm.utils import ApexScaler, NativeScaler
 
-import tlmm.models
-from tlmm.data import create_token_label_target, TokenLabelMixup, FastCollateTokenLabelMixup, create_token_label_loader, create_token_label_dataset
-from tlmm.loss import TokenLabelCrossEntropy, TokenLabelSoftTargetCrossEntropy
-from tlmm.utils import load_for_transfer_learning
+import tlt.models
+from tlt.data import create_token_label_target, TokenLabelMixup, FastCollateTokenLabelMixup, create_token_label_loader, create_token_label_dataset
+from tlt.loss import TokenLabelCrossEntropy, TokenLabelSoftTargetCrossEntropy
+from tlt.utils import load_pretrained_weights
 
 
 try:
@@ -359,7 +359,7 @@ def main():
         args.num_classes = model.num_classes  # FIXME handle model default vs config num_classes more elegantly
 
     if args.finetune:
-        load_for_transfer_learning(model=model,checkpoint_path=args.finetune,use_ema=args.model_ema, strict=False, num_classes=args.num_classes)
+        load_pretrained_weights(model=model,checkpoint_path=args.finetune,use_ema=args.model_ema, strict=False, num_classes=args.num_classes)
 
     if args.local_rank == 0:
         _logger.info('Model %s created, param count: %d' %

diff --git a/setup.py b/setup.py
@@ -1,11 +1,11 @@
 from setuptools import setup, find_packages
 
 setup(
-  name = 'tlmm',
+  name = 'tlt',
   packages = find_packages(exclude=['seg','visualize']),
   version = '0.1.0',
   license='Apache License 2.0',
-  description = 'Token labeling for training image models',
+  description = 'Token Labeling Toolbox for training image models',
   author = 'Zihang Jiang',
   author_email = 'jzh0103@gmail.com',
   url = 'https://github.com/zihangJiang/TokenLabeling',

diff --git a/tlmm/utils/__init__.py b/tlmm/utils/__init__.py
diff --git a/tlmm/__init__.py → tlt/__init__.py b/tlmm/__init__.py → tlt/__init__.py
diff --git a/tlmm/data/__init__.py → tlt/data/__init__.py b/tlmm/data/__init__.py → tlt/data/__init__.py
diff --git a/tlmm/data/dataset.py → tlt/data/dataset.py b/tlmm/data/dataset.py → tlt/data/dataset.py
diff --git a/tlmm/data/label_transforms_factory.py → tlt/data/label_transforms_factory.py b/tlmm/data/label_transforms_factory.py → tlt/data/label_transforms_factory.py
diff --git a/tlmm/data/loader.py → tlt/data/loader.py b/tlmm/data/loader.py → tlt/data/loader.py
diff --git a/tlmm/data/mixup.py → tlt/data/mixup.py b/tlmm/data/mixup.py → tlt/data/mixup.py
diff --git a/tlmm/data/random_augment_label.py → tlt/data/random_augment_label.py b/tlmm/data/random_augment_label.py → tlt/data/random_augment_label.py
diff --git a/tlmm/loss/__init__.py → tlt/loss/__init__.py b/tlmm/loss/__init__.py → tlt/loss/__init__.py
diff --git a/tlmm/loss/cross_entropy.py → tlt/loss/cross_entropy.py b/tlmm/loss/cross_entropy.py → tlt/loss/cross_entropy.py
diff --git a/tlmm/models/__init__.py → tlt/models/__init__.py b/tlmm/models/__init__.py → tlt/models/__init__.py
diff --git a/tlmm/models/layers.py → tlt/models/layers.py b/tlmm/models/layers.py → tlt/models/layers.py
diff --git a/tlmm/models/lvvit.py → tlt/models/lvvit.py b/tlmm/models/lvvit.py → tlt/models/lvvit.py
diff --git a/tlt/utils/__init__.py b/tlt/utils/__init__.py
@@ -0,0 +1 @@
+from .utils import load_pretrained_weights
diff --git a/tlmm/utils/utils.py → tlt/utils/utils.py b/tlmm/utils/utils.py → tlt/utils/utils.py
@@ -106,7 +106,7 @@ def load_state_dict(checkpoint_path,model, use_ema=False, num_classes=1000):
         raise FileNotFoundError()
 
 
-def load_for_transfer_learning(model, checkpoint_path, use_ema=False, strict=True, num_classes=1000):
+def load_pretrained_weights(model, checkpoint_path, use_ema=False, strict=True, num_classes=1000):
     state_dict = load_state_dict(checkpoint_path, model, use_ema, num_classes)
     model.load_state_dict(state_dict, strict=strict)
 

diff --git a/validate.py b/validate.py
@@ -20,7 +20,7 @@
 from timm.models.helpers import load_state_dict
 from timm.data import create_dataset, create_loader, resolve_data_config, RealLabelsImagenet
 from timm.utils import accuracy, AverageMeter, natural_key, setup_default_logging, set_jit_legacy
-import tlmm.models
+import tlt.models
 
 has_apex = False
 try: