opendet3d/zoo/gdino3d/base/optim.py

"""3D Grounding DINO Optimizer config."""

from __future__ import annotations

from ml_collections import ConfigDict
from torch.optim.adamw import AdamW
from torch.optim.lr_scheduler import LinearLR, MultiStepLR
from vis4d.config import class_config
from vis4d.config.typing import ExperimentParameters
from vis4d.zoo.base import get_lr_scheduler_cfg, get_optimizer_cfg


def get_optim_cfg(
    params: ExperimentParameters,
    epoch_based: bool = True,
    convnext: bool = False,
    warmup: bool = True,
    param_groups: list[dict] | None = None,
) -> list[ConfigDict]:
    """Returns the optimizer configuration."""
    if convnext:
        param_groups = get_convnext_param_groups()
    elif param_groups is None:
        param_groups = [
            {"custom_keys": ["language_model"], "lr_mult": 0.1},
            {"custom_keys": ["backbone"], "lr_mult": 0.1},
        ]

    lr_schedulers = []
    if warmup:
        lr_schedulers.append(
            get_lr_scheduler_cfg(
                class_config(LinearLR, start_factor=0.1, total_iters=1000),
                end=1000,
                epoch_based=False,
            )
        )

    lr_schedulers.append(
        get_lr_scheduler_cfg(
            class_config(
                MultiStepLR,
                milestones=[params.step_1, params.step_2],
                gamma=0.1,
            ),
            epoch_based=epoch_based,
        ),
    )

    optimizers = [
        get_optimizer_cfg(
            optimizer=class_config(
                AdamW, lr=params.lr, weight_decay=params.weight_decay
            ),
            lr_schedulers=lr_schedulers,
            param_groups=param_groups,
        )
    ]
    return optimizers


def get_convnext_param_groups(
    language_scale: float = 0.1, backbone_scale: float = 0.1
) -> list[dict[str : [list[str], float]]]:
    """Returns the parameter groups for ConvNeXt-Base."""
    return [
        {"custom_keys": ["language_model"], "lr_mult": language_scale},
        {
            "custom_keys": [
                "backbone.downsample_layers.0.0.weight",
                "backbone.downsample_layers.0.0.bias",
                "backbone.downsample_layers.0.1.weight",
                "backbone.downsample_layers.0.1.bias",
            ],
            "lr_mult": backbone_scale * 0.054975581388800036,
        },
        {
            "custom_keys": [
                "backbone.downsample_layers.1.0.weight",
                "backbone.downsample_layers.1.0.bias",
                "backbone.downsample_layers.1.1.weight",
                "backbone.downsample_layers.1.1.bias",
                "backbone.stages.1.0.gamma",
                "backbone.stages.1.0.dwconv.weight",
                "backbone.stages.1.0.dwconv.bias",
                "backbone.stages.1.0.norm.weight",
                "backbone.stages.1.0.norm.bias",
                "backbone.stages.1.0.pwconv1.weight",
                "backbone.stages.1.0.pwconv1.bias",
                "backbone.stages.1.0.pwconv2.weight",
                "backbone.stages.1.0.pwconv2.bias",
                "backbone.stages.1.1.gamma",
                "backbone.stages.1.1.dwconv.weight",
                "backbone.stages.1.1.dwconv.bias",
                "backbone.stages.1.1.norm.weight",
                "backbone.stages.1.1.norm.bias",
                "backbone.stages.1.1.pwconv1.weight",
                "backbone.stages.1.1.pwconv1.bias",
                "backbone.stages.1.1.pwconv2.weight",
                "backbone.stages.1.1.pwconv2.bias",
                "backbone.stages.1.2.gamma",
                "backbone.stages.1.2.dwconv.weight",
                "backbone.stages.1.2.dwconv.bias",
                "backbone.stages.1.2.norm.weight",
                "backbone.stages.1.2.norm.bias",
                "backbone.stages.1.2.pwconv1.weight",
                "backbone.stages.1.2.pwconv1.bias",
                "backbone.stages.1.2.pwconv2.weight",
                "backbone.stages.1.2.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.08589934592000005,
        },
        {
            "custom_keys": [
                "backbone.downsample_layers.2.0.weight",
                "backbone.downsample_layers.2.0.bias",
                "backbone.downsample_layers.2.1.weight",
                "backbone.downsample_layers.2.1.bias",
                "backbone.stages.2.0.gamma",
                "backbone.stages.2.0.dwconv.weight",
                "backbone.stages.2.0.dwconv.bias",
                "backbone.stages.2.0.norm.weight",
                "backbone.stages.2.0.norm.bias",
                "backbone.stages.2.0.pwconv1.weight",
                "backbone.stages.2.0.pwconv1.bias",
                "backbone.stages.2.0.pwconv2.weight",
                "backbone.stages.2.0.pwconv2.bias",
                "backbone.stages.2.1.gamma",
                "backbone.stages.2.1.dwconv.weight",
                "backbone.stages.2.1.dwconv.bias",
                "backbone.stages.2.1.norm.weight",
                "backbone.stages.2.1.norm.bias",
                "backbone.stages.2.1.pwconv1.weight",
                "backbone.stages.2.1.pwconv1.bias",
                "backbone.stages.2.1.pwconv2.weight",
                "backbone.stages.2.1.pwconv2.bias",
                "backbone.stages.2.2.gamma",
                "backbone.stages.2.2.dwconv.weight",
                "backbone.stages.2.2.dwconv.bias",
                "backbone.stages.2.2.norm.weight",
                "backbone.stages.2.2.norm.bias",
                "backbone.stages.2.2.pwconv1.weight",
                "backbone.stages.2.2.pwconv1.bias",
                "backbone.stages.2.2.pwconv2.weight",
                "backbone.stages.2.2.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.10737418240000006,
        },
        {
            "custom_keys": [
                "backbone.downsample_layers.3.0.weight",
                "backbone.downsample_layers.3.0.bias",
                "backbone.downsample_layers.3.1.weight",
                "backbone.downsample_layers.3.1.bias",
                "backbone.stages.3.0.gamma",
                "backbone.stages.3.0.dwconv.weight",
                "backbone.stages.3.0.dwconv.bias",
                "backbone.stages.3.0.norm.weight",
                "backbone.stages.3.0.norm.bias",
                "backbone.stages.3.0.pwconv1.weight",
                "backbone.stages.3.0.pwconv1.bias",
                "backbone.stages.3.0.pwconv2.weight",
                "backbone.stages.3.0.pwconv2.bias",
                "backbone.stages.3.1.gamma",
                "backbone.stages.3.1.dwconv.weight",
                "backbone.stages.3.1.dwconv.bias",
                "backbone.stages.3.1.norm.weight",
                "backbone.stages.3.1.norm.bias",
                "backbone.stages.3.1.pwconv1.weight",
                "backbone.stages.3.1.pwconv1.bias",
                "backbone.stages.3.1.pwconv2.weight",
                "backbone.stages.3.1.pwconv2.bias",
                "backbone.stages.3.2.gamma",
                "backbone.stages.3.2.dwconv.weight",
                "backbone.stages.3.2.dwconv.bias",
                "backbone.stages.3.2.norm.weight",
                "backbone.stages.3.2.norm.bias",
                "backbone.stages.3.2.pwconv1.weight",
                "backbone.stages.3.2.pwconv1.bias",
                "backbone.stages.3.2.pwconv2.weight",
                "backbone.stages.3.2.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 1.25,
        },
        {
            "custom_keys": [
                "backbone.stages.0.0.gamma",
                "backbone.stages.0.0.dwconv.weight",
                "backbone.stages.0.0.dwconv.bias",
                "backbone.stages.0.0.norm.weight",
                "backbone.stages.0.0.norm.bias",
                "backbone.stages.0.0.pwconv1.weight",
                "backbone.stages.0.0.pwconv1.bias",
                "backbone.stages.0.0.pwconv2.weight",
                "backbone.stages.0.0.pwconv2.bias",
                "backbone.stages.0.1.gamma",
                "backbone.stages.0.1.dwconv.weight",
                "backbone.stages.0.1.dwconv.bias",
                "backbone.stages.0.1.norm.weight",
                "backbone.stages.0.1.norm.bias",
                "backbone.stages.0.1.pwconv1.weight",
                "backbone.stages.0.1.pwconv1.bias",
                "backbone.stages.0.1.pwconv2.weight",
                "backbone.stages.0.1.pwconv2.bias",
                "backbone.stages.0.2.gamma",
                "backbone.stages.0.2.dwconv.weight",
                "backbone.stages.0.2.dwconv.bias",
                "backbone.stages.0.2.norm.weight",
                "backbone.stages.0.2.norm.bias",
                "backbone.stages.0.2.pwconv1.weight",
                "backbone.stages.0.2.pwconv1.bias",
                "backbone.stages.0.2.pwconv2.weight",
                "backbone.stages.0.2.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.06871947673600004,
        },
        {
            "custom_keys": [
                "backbone.stages.2.3.gamma",
                "backbone.stages.2.3.dwconv.weight",
                "backbone.stages.2.3.dwconv.bias",
                "backbone.stages.2.3.norm.weight",
                "backbone.stages.2.3.norm.bias",
                "backbone.stages.2.3.pwconv1.weight",
                "backbone.stages.2.3.pwconv1.bias",
                "backbone.stages.2.3.pwconv2.weight",
                "backbone.stages.2.3.pwconv2.bias",
                "backbone.stages.2.4.gamma",
                "backbone.stages.2.4.dwconv.weight",
                "backbone.stages.2.4.dwconv.bias",
                "backbone.stages.2.4.norm.weight",
                "backbone.stages.2.4.norm.bias",
                "backbone.stages.2.4.pwconv1.weight",
                "backbone.stages.2.4.pwconv1.bias",
                "backbone.stages.2.4.pwconv2.weight",
                "backbone.stages.2.4.pwconv2.bias",
                "backbone.stages.2.5.gamma",
                "backbone.stages.2.5.dwconv.weight",
                "backbone.stages.2.5.dwconv.bias",
                "backbone.stages.2.5.norm.weight",
                "backbone.stages.2.5.norm.bias",
                "backbone.stages.2.5.pwconv1.weight",
                "backbone.stages.2.5.pwconv1.bias",
                "backbone.stages.2.5.pwconv2.weight",
                "backbone.stages.2.5.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.13421772800000006,
        },
        {
            "custom_keys": [
                "backbone.stages.2.6.gamma",
                "backbone.stages.2.6.dwconv.weight",
                "backbone.stages.2.6.dwconv.bias",
                "backbone.stages.2.6.norm.weight",
                "backbone.stages.2.6.norm.bias",
                "backbone.stages.2.6.pwconv1.weight",
                "backbone.stages.2.6.pwconv1.bias",
                "backbone.stages.2.6.pwconv2.weight",
                "backbone.stages.2.6.pwconv2.bias",
                "backbone.stages.2.7.gamma",
                "backbone.stages.2.7.dwconv.weight",
                "backbone.stages.2.7.dwconv.bias",
                "backbone.stages.2.7.norm.weight",
                "backbone.stages.2.7.norm.bias",
                "backbone.stages.2.7.pwconv1.weight",
                "backbone.stages.2.7.pwconv1.bias",
                "backbone.stages.2.7.pwconv2.weight",
                "backbone.stages.2.7.pwconv2.bias",
                "backbone.stages.2.8.gamma",
                "backbone.stages.2.8.dwconv.weight",
                "backbone.stages.2.8.dwconv.bias",
                "backbone.stages.2.8.norm.weight",
                "backbone.stages.2.8.norm.bias",
                "backbone.stages.2.8.pwconv1.weight",
                "backbone.stages.2.8.pwconv1.bias",
                "backbone.stages.2.8.pwconv2.weight",
                "backbone.stages.2.8.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.1677721600000001,
        },
        {
            "custom_keys": [
                "backbone.stages.2.9.gamma",
                "backbone.stages.2.9.dwconv.weight",
                "backbone.stages.2.9.dwconv.bias",
                "backbone.stages.2.9.norm.weight",
                "backbone.stages.2.9.norm.bias",
                "backbone.stages.2.9.pwconv1.weight",
                "backbone.stages.2.9.pwconv1.bias",
                "backbone.stages.2.9.pwconv2.weight",
                "backbone.stages.2.9.pwconv2.bias",
                "backbone.stages.2.10.gamma",
                "backbone.stages.2.10.dwconv.weight",
                "backbone.stages.2.10.dwconv.bias",
                "backbone.stages.2.10.norm.weight",
                "backbone.stages.2.10.norm.bias",
                "backbone.stages.2.10.pwconv1.weight",
                "backbone.stages.2.10.pwconv1.bias",
                "backbone.stages.2.10.pwconv2.weight",
                "backbone.stages.2.10.pwconv2.bias",
                "backbone.stages.2.11.gamma",
                "backbone.stages.2.11.dwconv.weight",
                "backbone.stages.2.11.dwconv.bias",
                "backbone.stages.2.11.norm.weight",
                "backbone.stages.2.11.norm.bias",
                "backbone.stages.2.11.pwconv1.weight",
                "backbone.stages.2.11.pwconv1.bias",
                "backbone.stages.2.11.pwconv2.weight",
                "backbone.stages.2.11.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.20971520000000007,
        },
        {
            "custom_keys": [
                "backbone.stages.2.12.gamma",
                "backbone.stages.2.12.dwconv.weight",
                "backbone.stages.2.12.dwconv.bias",
                "backbone.stages.2.12.norm.weight",
                "backbone.stages.2.12.norm.bias",
                "backbone.stages.2.12.pwconv1.weight",
                "backbone.stages.2.12.pwconv1.bias",
                "backbone.stages.2.12.pwconv2.weight",
                "backbone.stages.2.12.pwconv2.bias",
                "backbone.stages.2.13.gamma",
                "backbone.stages.2.13.dwconv.weight",
                "backbone.stages.2.13.dwconv.bias",
                "backbone.stages.2.13.norm.weight",
                "backbone.stages.2.13.norm.bias",
                "backbone.stages.2.13.pwconv1.weight",
                "backbone.stages.2.13.pwconv1.bias",
                "backbone.stages.2.13.pwconv2.weight",
                "backbone.stages.2.13.pwconv2.bias",
                "backbone.stages.2.14.gamma",
                "backbone.stages.2.14.dwconv.weight",
                "backbone.stages.2.14.dwconv.bias",
                "backbone.stages.2.14.norm.weight",
                "backbone.stages.2.14.norm.bias",
                "backbone.stages.2.14.pwconv1.weight",
                "backbone.stages.2.14.pwconv1.bias",
                "backbone.stages.2.14.pwconv2.weight",
                "backbone.stages.2.14.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.2621440000000001,
        },
        {
            "custom_keys": [
                "backbone.stages.2.15.gamma",
                "backbone.stages.2.15.dwconv.weight",
                "backbone.stages.2.15.dwconv.bias",
                "backbone.stages.2.15.norm.weight",
                "backbone.stages.2.15.norm.bias",
                "backbone.stages.2.15.pwconv1.weight",
                "backbone.stages.2.15.pwconv1.bias",
                "backbone.stages.2.15.pwconv2.weight",
                "backbone.stages.2.15.pwconv2.bias",
                "backbone.stages.2.16.gamma",
                "backbone.stages.2.16.dwconv.weight",
                "backbone.stages.2.16.dwconv.bias",
                "backbone.stages.2.16.norm.weight",
                "backbone.stages.2.16.norm.bias",
                "backbone.stages.2.16.pwconv1.weight",
                "backbone.stages.2.16.pwconv1.bias",
                "backbone.stages.2.16.pwconv2.weight",
                "backbone.stages.2.16.pwconv2.bias",
                "backbone.stages.2.17.gamma",
                "backbone.stages.2.17.dwconv.weight",
                "backbone.stages.2.17.dwconv.bias",
                "backbone.stages.2.17.norm.weight",
                "backbone.stages.2.17.norm.bias",
                "backbone.stages.2.17.pwconv1.weight",
                "backbone.stages.2.17.pwconv1.bias",
                "backbone.stages.2.17.pwconv2.weight",
                "backbone.stages.2.17.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.3276800000000001,
        },
        {
            "custom_keys": [
                "backbone.stages.2.18.gamma",
                "backbone.stages.2.18.dwconv.weight",
                "backbone.stages.2.18.dwconv.bias",
                "backbone.stages.2.18.norm.weight",
                "backbone.stages.2.18.norm.bias",
                "backbone.stages.2.18.pwconv1.weight",
                "backbone.stages.2.18.pwconv1.bias",
                "backbone.stages.2.18.pwconv2.weight",
                "backbone.stages.2.18.pwconv2.bias",
                "backbone.stages.2.19.gamma",
                "backbone.stages.2.19.dwconv.weight",
                "backbone.stages.2.19.dwconv.bias",
                "backbone.stages.2.19.norm.weight",
                "backbone.stages.2.19.norm.bias",
                "backbone.stages.2.19.pwconv1.weight",
                "backbone.stages.2.19.pwconv1.bias",
                "backbone.stages.2.19.pwconv2.weight",
                "backbone.stages.2.19.pwconv2.bias",
                "backbone.stages.2.20.gamma",
                "backbone.stages.2.20.dwconv.weight",
                "backbone.stages.2.20.dwconv.bias",
                "backbone.stages.2.20.norm.weight",
                "backbone.stages.2.20.norm.bias",
                "backbone.stages.2.20.pwconv1.weight",
                "backbone.stages.2.20.pwconv1.bias",
                "backbone.stages.2.20.pwconv2.weight",
                "backbone.stages.2.20.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.4096000000000001,
        },
        {
            "custom_keys": [
                "backbone.stages.2.21.gamma",
                "backbone.stages.2.21.dwconv.weight",
                "backbone.stages.2.21.dwconv.bias",
                "backbone.stages.2.21.norm.weight",
                "backbone.stages.2.21.norm.bias",
                "backbone.stages.2.21.pwconv1.weight",
                "backbone.stages.2.21.pwconv1.bias",
                "backbone.stages.2.21.pwconv2.weight",
                "backbone.stages.2.21.pwconv2.bias",
                "backbone.stages.2.22.gamma",
                "backbone.stages.2.22.dwconv.weight",
                "backbone.stages.2.22.dwconv.bias",
                "backbone.stages.2.22.norm.weight",
                "backbone.stages.2.22.norm.bias",
                "backbone.stages.2.22.pwconv1.weight",
                "backbone.stages.2.22.pwconv1.bias",
                "backbone.stages.2.22.pwconv2.weight",
                "backbone.stages.2.22.pwconv2.bias",
                "backbone.stages.2.23.gamma",
                "backbone.stages.2.23.dwconv.weight",
                "backbone.stages.2.23.dwconv.bias",
                "backbone.stages.2.23.norm.weight",
                "backbone.stages.2.23.norm.bias",
                "backbone.stages.2.23.pwconv1.weight",
                "backbone.stages.2.23.pwconv1.bias",
                "backbone.stages.2.23.pwconv2.weight",
                "backbone.stages.2.23.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.5120000000000001,
        },
        {
            "custom_keys": [
                "backbone.stages.2.24.gamma",
                "backbone.stages.2.24.dwconv.weight",
                "backbone.stages.2.24.dwconv.bias",
                "backbone.stages.2.24.norm.weight",
                "backbone.stages.2.24.norm.bias",
                "backbone.stages.2.24.pwconv1.weight",
                "backbone.stages.2.24.pwconv1.bias",
                "backbone.stages.2.24.pwconv2.weight",
                "backbone.stages.2.24.pwconv2.bias",
                "backbone.stages.2.25.gamma",
                "backbone.stages.2.25.dwconv.weight",
                "backbone.stages.2.25.dwconv.bias",
                "backbone.stages.2.25.norm.weight",
                "backbone.stages.2.25.norm.bias",
                "backbone.stages.2.25.pwconv1.weight",
                "backbone.stages.2.25.pwconv1.bias",
                "backbone.stages.2.25.pwconv2.weight",
                "backbone.stages.2.25.pwconv2.bias",
                "backbone.stages.2.26.gamma",
                "backbone.stages.2.26.dwconv.weight",
                "backbone.stages.2.26.dwconv.bias",
                "backbone.stages.2.26.norm.weight",
                "backbone.stages.2.26.norm.bias",
                "backbone.stages.2.26.pwconv1.weight",
                "backbone.stages.2.26.pwconv1.bias",
                "backbone.stages.2.26.pwconv2.weight",
                "backbone.stages.2.26.pwconv2.bias",
            ],
            "lr_mult": backbone_scale * 0.6400000000000001,
        },
    ]