support layer decay and different lr for text/visual encoder

src/open_clip/hf_model.py

@@ -162,3 +162,8 @@ def set_grad_checkpointing(self, enable=True):
 
     def init_parameters(self):
         pass
+
+    def get_num_layers(self):
+        encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
+        layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
+        return len(layer_list)

src/open_clip/transformer.py

@@ -374,6 +374,13 @@ def init_parameters(self):
     def set_grad_checkpointing(self, enable=True):
         self.transformer.grad_checkpointing = enable
 
+    def get_num_layers(self):
+        return self.transformer.layers
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'positional_embedding', 'class_embedding'}
+
     def forward(self, x: torch.Tensor):
         x = self.conv1(x)  # shape = [*, width, grid, grid]
         x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
@@ -461,6 +468,13 @@ def init_parameters(self):
     def set_grad_checkpointing(self, enable=True):
         self.transformer.grad_checkpointing = enable
 
+    def get_num_layers(self):
+        return self.transformer.layers
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'positional_embedding'}
+
     def build_attention_mask(self):
         # lazily create causal attention mask, with full attention between the vision tokens
         # pytorch uses additive attention mask; fill with -inf

src/training/main.py

@@ -31,7 +31,7 @@
 from training.params import parse_args
 from training.scheduler import cosine_lr
 from training.train import train_one_epoch, evaluate
-
+from training.optim_factory import create_optimizer
 
 def random_seed(seed=42, rank=0):
     torch.manual_seed(seed + rank)
@@ -177,22 +177,8 @@ def main(args):
     if args.train_data or args.dataset_type == "synthetic":
         assert not args.trace, 'Cannot train with traced model'
 
-        exclude = lambda n, p: p.ndim < 2 or "bn" in n or "ln" in n or "bias" in n or 'logit_scale' in n
-        include = lambda n, p: not exclude(n, p)
-
-        named_parameters = list(model.named_parameters())
-        gain_or_bias_params = [p for n, p in named_parameters if exclude(n, p) and p.requires_grad]
-        rest_params = [p for n, p in named_parameters if include(n, p) and p.requires_grad]
-
-        optimizer = optim.AdamW(
-            [
-                {"params": gain_or_bias_params, "weight_decay": 0.},
-                {"params": rest_params, "weight_decay": args.wd},
-            ],
-            lr=args.lr,
-            betas=(args.beta1, args.beta2),
-            eps=args.eps,
-        )
+        optimizer = create_optimizer(args, model)
+
         if args.horovod:
             optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
             hvd.broadcast_parameters(model.state_dict(), root_rank=0)
@@ -232,7 +218,7 @@ def main(args):
     scheduler = None
     if 'train' in data and optimizer is not None:
         total_steps = (data["train"].dataloader.num_batches // args.accum_freq) * args.epochs
-        scheduler = cosine_lr(optimizer, args.lr, args.warmup, total_steps)
+        scheduler = cosine_lr(optimizer, args, total_steps)
 
     # determine if this worker should save logs and checkpoints. only do so if it is rank == 0
     args.save_logs = args.logs and args.logs.lower() != 'none' and is_master(args)

src/training/optim_factory.py

@@ -0,0 +1,185 @@
+"""
+Adapted from BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py
+"""
+import logging
+from torch import optim
+
+
+def get_num_layer_for_transformer(var_name, num_max_layer):
+    first_layer_var_names = [
+        "visual.cls_token", 
+        "visual.mask_token", 
+        "visual.pos_embed", 
+        "visual.positional_embedding", 
+        "visual.patch_embed",
+        "visual.conv1", # name of patch embed in CLIP
+        "text.pos_embed", 
+        "text.positional_embedding",
+        "text.token_embedding",
+        "text.transformer.embeddings.word_embeddings",
+        "text.transformer.embeddings.position_embeddings",
+        "text.transformer.embeddings.token_type_embeddings"
+    ]
+    for first_layer_var_name in first_layer_var_names:
+        if var_name.startswith(first_layer_var_name):
+            return 0
+
+    if var_name.startswith("visual.rel_pos_bias"):
+        layer_id = num_max_layer - 1
+    elif var_name.startswith("visual.blocks"):
+        layer_id = int(var_name.split('.')[2]) + 1
+    elif var_name.startswith("visual.transformer.resblocks"):
+        layer_id = int(var_name.split('.')[3]) + 1
+    elif var_name.startswith("text.transformer.resblocks"):
+        layer_id = int(var_name.split('.')[3]) + 1
+    elif var_name.startswith("text.transformer.encoder.layer"):
+        layer_id = int(var_name.split('.')[4]) + 1
+    else:
+        logging.warning(f'Unknown layer for {var_name} when setting layer decay. Defaulting to {num_max_layer-1}.')
+        layer_id = num_max_layer - 1
+
+    return layer_id
+
+
+class LayerDecayValueAssigner(object):
+    def __init__(self, layer_decay, num_layers):
+        self.values = [
+            layer_decay ** (num_layers + 1 - i)
+            for i in range(num_layers + 2)
+        ]
+
+    def get_scale(self, layer_id):
+        return self.values[layer_id]
+
+    def get_layer_id(self, var_name):
+        return get_num_layer_for_transformer(var_name, len(self.values))
+
+
+def do_weight_decay(name, param):
+    return (
+        param.ndim <= 1 or "bn" in name or
+        "ln" in name or
+        "bias" in name or
+        "logit_scale" in name
+    )
+
+
+def param_groups_layer_decay(model_params, lr, weight_decay, lr_scale_assigner, tower):
+    param_group_vars = {}
+    for name, param in model_params:
+        if not param.requires_grad:
+            continue
+
+        if do_weight_decay(name, param):
+            group_name = "decay"
+            this_weight_decay = weight_decay
+        else:
+            group_name = "no_decay"
+            this_weight_decay = 0.
+
+        if lr_scale_assigner:
+            layer_id = lr_scale_assigner.get_layer_id(name)
+            scale = lr_scale_assigner.get_layer_scale(layer_id)
+            group_name = f"{tower}_layer_{layer_id}_{group_name}"
+        else:
+            layer_id = None
+            scale = 1.0
+
+        if group_name not in param_group_vars:
+            param_group_vars[group_name] = {
+                "group": tower,
+                "weight_decay": this_weight_decay,
+                "params": [],
+                "lr_scale": scale,
+                "lr": lr
+            }
+
+        param_group_vars[group_name]["params"].append(param)
+    return list(param_group_vars.values())
+
+
+def get_all_param_groups(args, model, lr_scale_assigner_visual, lr_scale_assigner_text):
+    visual_params, text_params, other_params = [], [], []
+    for name, param in model.named_parameters():
+        if name.startswith('visual.'):
+            visual_params.append([name, param])
+        elif name.startswith('text.'):
+            text_params.append([name, param])
+        else:
+            if lr_scale_assigner_visual is not None or lr_scale_assigner_text is not None:
+                logging.warning(f"Param {param} is not assigned to either the visual or text encoders.")
+            other_params.append([name, param])
+
+    visual_optim_params = param_groups_layer_decay(
+        visual_params, 
+        args.visual_lr or args.lr, 
+        args.visual_wd or args.wd, 
+        lr_scale_assigner_visual, 
+        'visual'
+    )
+    text_optim_params = param_groups_layer_decay(
+        text_params, 
+        args.text_lr or args.lr, 
+        args.text_wd or args.wd, 
+        lr_scale_assigner_text, 
+        'text'
+    )
+    other_optim_params = param_groups_layer_decay(
+        other_params, 
+        args.lr, 
+        args.wd, 
+        None, 
+        'other'
+    )
+
+    optim_params = visual_optim_params + text_optim_params + other_optim_params
+
+    if len(optim_params) == 0:
+        optim_params = model.parameters()
+    return optim_params
+
+
+def check_lr_layer_decay_available(model):
+    def is_available(model, white_list):
+        for name, _ in model.named_parameters():
+            for white_name in white_list:
+                if name.startswith(white_name):
+                    return True
+        return False
+
+    visual_white_list = ["visual.blocks", "visual.transformer.resblocks"]
+    text_white_list = ["text.transformer.resblocks", "text.transformer.encoder.layer"]
+
+    do_visual_lr_decay = is_available(model, visual_white_list)
+    do_text_lr_decay = is_available(model, text_white_list)
+
+    if not do_visual_lr_decay:
+        logging.warning("Learning rate layer decay is currently only supported for built-in ViT models.")
+    if not do_text_lr_decay:
+        logging.warning("Learning rate layer decay is currently only supported for built-in text transformers.")
+    return do_visual_lr_decay, do_text_lr_decay
+
+
+def create_optimizer(args, model):
+    lr_scale_assigner_visual, lr_scale_assigner_text = None, None
+    do_visual_lr_decay, do_text_lr_decay = check_lr_layer_decay_available(model)
+
+    if args.visual_ld and do_visual_lr_decay:
+        visual_num_layers = model.visual.get_num_layers()
+        lr_scale_assigner_visual = LayerDecayValueAssigner(args.visual_ld, visual_num_layers)
+        logging.info("Assigned visual lr scale values = %s" % str(lr_scale_assigner_visual.values))
+
+    if args.text_ld and do_text_lr_decay:
+        text_num_layers = model.text.get_num_layers()
+        lr_scale_assigner_text = LayerDecayValueAssigner(args.text_ld, text_num_layers)
+        logging.info("Assigned text lr scale values = %s" % str(lr_scale_assigner_text.values))
+
+    optim_params = get_all_param_groups(args, model, lr_scale_assigner_visual, lr_scale_assigner_text)
+
+    optim_args = dict(
+        betas=(args.beta1, args.beta2),
+        eps=args.eps,
+    )
+
+    optimizer = optim.AdamW(optim_params, **optim_args)
+    return optimizer

src/training/params.py

@@ -106,10 +106,17 @@ def parse_args(args):
         "--epochs", type=int, default=32, help="Number of epochs to train for."
     )
     parser.add_argument("--lr", type=float, default=None, help="Learning rate.")
+    parser.add_argument("--text-lr", type=float, default=None, help="Learning rate of text tower.")
+    parser.add_argument("--visual-lr", type=float, default=None, help="Learning rate of image tower.")
     parser.add_argument("--beta1", type=float, default=None, help="Adam beta 1.")
     parser.add_argument("--beta2", type=float, default=None, help="Adam beta 2.")
     parser.add_argument("--eps", type=float, default=None, help="Adam epsilon.")
     parser.add_argument("--wd", type=float, default=0.2, help="Weight decay.")
+    parser.add_argument("--text-wd", type=float, default=None, help="Weight decay of text tower.")
+    parser.add_argument("--visual-wd", type=float, default=None, help="Weight decay of image tower.")
+    parser.add_argument("--ld", type=float, default=1.0, help="Learning rate Layer decay.")
+    parser.add_argument("--text-ld", type=float, default=None, help="Learning rate Layer decay of text tower.")
+    parser.add_argument("--visual-ld", type=float, default=None, help="Learning rate Layer decay of image tower.")
     parser.add_argument(
         "--warmup", type=int, default=10000, help="Number of steps to warmup for."
     )

src/training/scheduler.py

@@ -10,14 +10,23 @@ def _warmup_lr(base_lr, warmup_length, step):
     return base_lr * (step + 1) / warmup_length
 
 
-def cosine_lr(optimizer, base_lr, warmup_length, steps):
+def cosine_lr(optimizer, args, steps):
     def _lr_adjuster(step):
-        if step < warmup_length:
-            lr = _warmup_lr(base_lr, warmup_length, step)
-        else:
-            e = step - warmup_length
-            es = steps - warmup_length
-            lr = 0.5 * (1 + np.cos(np.pi * e / es)) * base_lr
-        assign_learning_rate(optimizer, lr)
+        for param_group in optimizer.param_groups:
+            if param_group['group'] == 'text':
+                base_lr = args.text_lr if args.text_lr is not None else args.lr
+            elif param_group['group'] == 'visual':
+                base_lr = args.visual_lr if args.visual_lr is not None else args.lr
+            else:
+                base_lr = args.lr
+
+            if step < args.warmup:
+                lr = _warmup_lr(base_lr, args.warmup, step)
+            else:
+                e = step - args.warmup
+                es = steps - args.warmup
+                lr = 0.5 * (1 + np.cos(np.pi * e / es)) * base_lr
+            scale = param_group.get("lr_scale", 1.0)
+            param_group["lr"] = scale * lr
         return lr
     return _lr_adjuster