libai.optim

libai.optim.build_optimizer(cfg, model)[source]

Build an optimizer from config.

libai.optim.get_default_optimizer_params(model, base_lr=None, weight_decay=None, weight_decay_norm=None, weight_decay_bias=None, clip_grad_max_norm=None, clip_grad_norm_type=None, overrides=None)[source]

Get default param list for optimizer, with suport for a few types of overrides. If no overrides are needed, it is equivalent to model.parameters().

Parameters
  • base_lr – lr for every group by default. Can be omitted to use the one in optimizer.

  • weight_decay – weight decay for every group by default. Can be omitted to use the one in optimizer.

  • weight_decay_norm – override weight decay for params in normalization layers

  • weight_decay_bias – override weight decay for bias parameters

  • overrides – if not None, provides values for optimizer hyperparameters (LR, weight decay) for module parameters with a given name; e.g. {"embedding": {"lr": 0.01, "weight_decay": 0.1}} will set the LR and weight decay values for all module parameters named embedding.

For common transformer models, weight_decay_norm and weight_decay_bias are usually set to 0.

Example:

flow.optim.AdamW(
    get_default_optimizer_params(model, weight_decay_norm=0, weight_decay_bias=0),
    lr=0.01,
    weight_decay=1e-4
)