EfficientNetV2B0 transfer learning on cifar10 testing freezing backbone

[leondgarse]

• Freezing backbone by set --freeze_backbone, will set layer.trainable=False till model GlobalAveragePooling2D.

  # First stage freezing backbone. total epochs = lr_decay_steps + lr_cooldown_steps = 22
  CUDA_VISIBLE_DEVICES='1' TF_XLA_FLAGS="--tf_xla_auto_jit=2" ./train_script.py -m efficientnet.EfficientNetV2B0 \
  --data_name cifar10 --lr_decay_steps 20 --lr_cooldown_steps 2 --batch_size 256 --pretrained imagenet \
  --freeze_backbone --summary --seed 0 -s freeze_backbone_test
  # [Search pool layer] header_layer_id = 277, layer.name = predictions
  # [Search pool layer] header_layer_id = 276, layer.name = head_drop
  # [Search pool layer] header_layer_id = 275, layer.name = avg_pool
  # ...
  # Trainable params: 12,810
  # Non-trainable params: 5,919,312

  Train a second stage using previous trained model with unfreezing backbone. Specify --pretrained checkpoints/{basic_save_name}_latest.h5 for reloading weights.

  # Second stage unfreezing backbone. Learning rate scheduler using cosine lr restart.
  # For default lr_t_mul=2, specify total epochs = lr_decay_steps * (1 + lr_t_mul) + lr_cooldown_steps * 2 = 64
  CUDA_VISIBLE_DEVICES='1' TF_XLA_FLAGS="--tf_xla_auto_jit=2" ./train_script.py -m efficientnet.EfficientNetV2B0 \
  --data_name cifar10 --lr_decay_steps 20 --lr_cooldown_steps 2 --batch_size 256 --epochs 64 \
  --pretrained checkpoints/freeze_backbone_test_latest.h5 --initial_epoch 22 --seed 0 -s freeze_backbone_test

• Without freezing backbone as a compare.

  CUDA_VISIBLE_DEVICES='1' TF_XLA_FLAGS="--tf_xla_auto_jit=2" ./train_script.py -m efficientnet.EfficientNetV2B0 \
  --data_name cifar10 --lr_decay_steps 20 --lr_cooldown_steps 2 --batch_size 256 --epochs 64 \
  --pretrained imagenet --seed 0 -s freeze_backbone_compare_test_lr_restart
  # ...
  # Trainable params: 5,871,514
  # Non-trainable params: 60,608

• Without freezing backbone, without learning rate restart as a compare.

  CUDA_VISIBLE_DEVICES='1' TF_XLA_FLAGS="--tf_xla_auto_jit=2" ./train_script.py -m efficientnet.EfficientNetV2B0 \
  --data_name cifar10 --lr_decay_steps 60 --lr_cooldown_steps 4 --batch_size 256 --epochs 64 \
  --pretrained imagenet --seed 0 -s freeze_backbone_compare_test_lr_not_restart

• Freeze BatchNormalization layers only by set freeze_norm_layers, will set layer.trainable=False for all BatchNormalization and LayerNormlization.

  # First stage freezing backbone. total epochs = lr_decay_steps + lr_cooldown_steps = 22
  CUDA_VISIBLE_DEVICES='1' TF_XLA_FLAGS="--tf_xla_auto_jit=2" ./train_script.py -m efficientnet.EfficientNetV2B0 \
  --data_name cifar10 --lr_decay_steps 20 --lr_cooldown_steps 2 --batch_size 256 --pretrained imagenet \
  --freeze_norm_layers --summary --seed 0 -s freeze_bn_test
  # ...
  # Trainable params: 5,810,906
  # Non-trainable params: 121,216

  Train a second stage using previous trained model with unfreezing norm layers. Specify --pretrained checkpoints/{basic_save_name}_latest.h5 for reloading weights.

  # Second stage unfreezing norm layers. Learning rate scheduler using cosine lr restart.
  # For default lr_t_mul=2, specify total epochs = lr_decay_steps * (1 + lr_t_mul) + lr_cooldown_steps * 2 = 64
  CUDA_VISIBLE_DEVICES='1' TF_XLA_FLAGS="--tf_xla_auto_jit=2" ./train_script.py -m efficientnet.EfficientNetV2B0 \
  --data_name cifar10 --lr_decay_steps 20 --lr_cooldown_steps 2 --batch_size 256 --epochs 64 \
  --pretrained checkpoints/freeze_bn_test_latest.h5 --initial_epoch 22 --seed 0 -s freeze_bn_test

• Freezing BatchNormalization layers for entire epochs

  CUDA_VISIBLE_DEVICES='1' TF_XLA_FLAGS="--tf_xla_auto_jit=2" ./train_script.py -m efficientnet.EfficientNetV2B0 \
  --data_name cifar10 --lr_decay_steps 20 --lr_cooldown_steps 2 --batch_size 256 --epochs 64 \
  --freeze_norm_layers --pretrained imagenet --seed 0 -s freeze_bn_test_entire
  # ...
  # Trainable params: 5,810,906
  # Non-trainable params: 121,216

• Results

  	Train acc, loss	Best val_acc, val_loss	Epoch 64 val_acc
  freeze backbone	0.78688, 0.1396	Epoch 56, 0.9792, 0.01478	0.9776
  w/o freeze backbone	0.79654, 0.1332	Epoch 60, 0.9759, 0.01479	0.9754
  w/o freeze backbone, w/o lr restart	0.79072, 0.1371	Epoch 54, 0.9719, 0.01694	0.9718
  freeze batchnorm only	0.79436, 0.1356	Epoch 60, 0.9757, 0.01571	0.9755
  freeze batchnorm only, entire epochs	0.79352, 0.1348	Epoch 60, 0.9752, 0.01722	0.9750

• Plot

  from keras_cv_attention_models.imagenet import eval_func
  hhs = {
      "effb0 cifar10, freeze backbone": "checkpoints/freeze_backbone_test_hist.json",
      "effb0 cifar10, w/o freeze backbone": "checkpoints/freeze_backbone_compare_test_lr_restart_hist.json",
      "effb0 cifar10, w/o freeze backbone, w/o lr restart": "checkpoints/freeze_backbone_compare_test_lr_not_restart_hist.json",
      "effb0 cifar10, freeze batchnorm only": "checkpoints/freeze_bn_test_hist.json",
      "effb0 cifar10, freeze batchnorm only, entire epochs": "checkpoints/freeze_bn_test_entire_hist.json",
  }
  
  fig = eval_func.plot_hists(hhs.values(), list(hhs.keys()), skip_first=3, base_size=8)