diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml index 8cf56b9a4..46f6886cf 100644 --- a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml @@ -152,6 +152,10 @@ robomimic_model: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] +# LR multiplier for the image encoders' pretrained backbones; 1.0 keeps the +# single optimizer param group (ModelWrapper.configure_optimizers). +backbone_lr_scale: 1.0 + optimizer: _target_: torch.optim.AdamW _partial_: true diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index 3a6d4e0a3..119b20590 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -586,6 +586,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ResNet(PolicyStem): + # ``self.net`` holds the ImageNet-pretrained torchvision backbone (``self.proj`` + # is built here and must keep its own init). Declared so a weight-protection + # pass can find it without knowing about this class. + _hpt_pretrained = True + _hpt_pretrained_attrs = ("net",) + def __init__( self, output_dim: int = 10, @@ -617,6 +623,7 @@ def __init__( self.avgpool = nn.AvgPool2d(7, stride=1) # Freeze the backbone if specified + self.freeze_backbone = freeze_backbone if freeze_backbone: self._freeze_backbone() @@ -630,6 +637,23 @@ def _freeze_backbone(self): for param in self.net.parameters(): param.requires_grad = False + def train(self, mode: bool = True): + """Keep a frozen backbone in eval mode regardless of the outer flag. + + ``requires_grad = False`` alone leaves the BatchNorm layers in train + mode: they would keep normalizing with batch statistics and keep + updating their running stats, so a "frozen" backbone would not be + frozen at all. + """ + super().train(mode) + if self.freeze_backbone: + self.net.eval() + return self + + def backbone_parameters(self): + """Pretrained-backbone parameters (``proj`` stays in the main group).""" + return list(self.net.parameters()) + def forward(self, x: torch.Tensor) -> torch.Tensor: """ Performs a forward pass of the model. diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index 4c73a04c1..662db1ddb 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -293,6 +293,33 @@ def on_validation_end(self): flush=True, ) + def _backbone_param_groups(self, lr: float, scale: float): + """Two param groups: the encoders' pretrained backbones at ``lr * + scale``, every other trainable parameter at ``lr``. + + A pretrained image backbone and a randomly initialised trunk do not + want the same learning rate, and a frozen-vs-tuned comparison under one + LR is confounded. Encoders that do not declare ``backbone_parameters`` + (and frozen parameters) are left to the main group. + """ + encoders = getattr(self.model.nets["policy"], "encoders", None) or {} + backbone, seen = [], set() + for encoder in encoders.values(): + if not hasattr(encoder, "backbone_parameters"): + continue + for param in encoder.backbone_parameters(): + if param.requires_grad and id(param) not in seen: + seen.add(id(param)) + backbone.append(param) + rest = [ + param + for param in self.trainer.model.parameters() + if param.requires_grad and id(param) not in seen + ] + groups = [{"params": backbone, "lr": lr * scale}, {"params": rest}] + # a fully frozen (or absent) backbone leaves the main group alone + return [group for group in groups if group["params"]] + def configure_optimizers(self) -> Dict[str, Any]: """Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you'd need one. But in the case of GANs or similar you might have multiple. @@ -305,12 +332,19 @@ def configure_optimizers(self) -> Dict[str, Any]: config_tree = getattr(self.hparams, "config_tree", None) if config_tree is not None: cfg = self._as_config(config_tree) - optimizer = hydra.utils.instantiate( - cfg.model.optimizer, - params=self.trainer.model.parameters(), - ) - if callable(optimizer): - optimizer = optimizer() + scale = float(cfg.model.get("backbone_lr_scale", 1.0)) + params = self.trainer.model.parameters() + if scale != 1.0: + lr = cfg.model.optimizer.get("lr") + if lr is None: + raise ValueError( + "model.backbone_lr_scale needs an explicit model.optimizer.lr" + ) + params = self._backbone_param_groups(float(lr), scale) + # Bound outside ``instantiate``, which would turn param-group dicts + # into config nodes. + factory = hydra.utils.instantiate(cfg.model.optimizer, _partial_=True) + optimizer = factory(params=params) scheduler_cfg = cfg.model.get("scheduler") if scheduler_cfg is not None: scheduler = hydra.utils.instantiate( diff --git a/tests/unit/test_vision_backbone.py b/tests/unit/test_vision_backbone.py new file mode 100644 index 000000000..6c284124a --- /dev/null +++ b/tests/unit/test_vision_backbone.py @@ -0,0 +1,199 @@ +"""The frozen-backbone eval-mode fix and the separate backbone LR group. + +Covers ``ResNet`` (unchanged forward, BatchNorm in eval when frozen, +``backbone_parameters``), ``ModelWrapper.configure_optimizers`` param groups, +and the flagship model config. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torchvision +from omegaconf import OmegaConf +from torch import nn + +from egomimic.models.hpt_nets import ResNet +from egomimic.pl_utils.pl_model import ModelWrapper + +# -------------------------------------------------------------------------- +# ResNet +# -------------------------------------------------------------------------- + + +def _resnet_stem(**kwargs): + torch.manual_seed(0) + return ResNet(output_dim=16, **kwargs) + + +def test_resnet_forward_matches_reference_implementation(): + """The baseline image path is bit-identical: the stem's output equals the + torchvision backbone followed by the stem's projection.""" + stem = _resnet_stem().eval() + x = torch.rand(2, 1, 1, 3, 96, 128) + with torch.no_grad(): + out = stem(x) + + backbone = nn.Sequential( + *list(torchvision.models.resnet18(weights="DEFAULT").children())[:-2] + ).eval() + with torch.no_grad(): + feat = backbone(x.view(-1, 3, 96, 128)) + expected = stem.proj(feat.view(2, feat.shape[1], -1).transpose(1, 2)) + assert out.shape == expected.shape + assert torch.equal(out, expected) + + +def test_resnet_frozen_backbone_batchnorm_stays_in_eval(): + """A frozen ResNet must not update / use batch statistics.""" + stem = _resnet_stem(freeze_backbone=True) + stem.train() + assert stem.training + bns = [ + m for m in stem.net.modules() if isinstance(m, nn.modules.batchnorm._BatchNorm) + ] + assert bns + assert all(not m.training for m in bns) + assert all(not p.requires_grad for p in stem.net.parameters()) + # an unfrozen stem keeps the default behaviour + trainable = _resnet_stem() + trainable.train() + assert all( + m.training + for m in trainable.net.modules() + if isinstance(m, nn.modules.batchnorm._BatchNorm) + ) + + +def test_resnet_frozen_backbone_running_stats_do_not_drift(): + stem = _resnet_stem(freeze_backbone=True) + stem.train() + bn = next( + m for m in stem.net.modules() if isinstance(m, nn.modules.batchnorm._BatchNorm) + ) + before = bn.running_mean.clone() + with torch.no_grad(): + stem(torch.rand(2, 1, 1, 3, 96, 128) * 5.0) + assert torch.equal(bn.running_mean, before) + + +def test_resnet_backbone_parameters_are_the_net(): + stem = _resnet_stem() + backbone = {id(p) for p in stem.backbone_parameters()} + assert backbone == {id(p) for p in stem.net.parameters()} + assert all(id(p) not in backbone for p in stem.proj.parameters()) + + +def test_resnet_backbone_parameters_with_multiple_copies(): + stem = ResNet(output_dim=16, num_of_copy=2) + backbone = {id(p) for p in stem.backbone_parameters()} + assert backbone == {id(p) for p in stem.net.parameters()} + + +# -------------------------------------------------------------------------- +# configure_optimizers +# -------------------------------------------------------------------------- + + +class _StubEncoder(nn.Module): + """Encoder with a pretrained-style backbone plus a trainable projection.""" + + def __init__(self, frozen: bool = False): + super().__init__() + self.net = nn.Linear(4, 4) + self.proj = nn.Linear(4, 4) + if frozen: + for p in self.net.parameters(): + p.requires_grad = False + + def backbone_parameters(self): + return list(self.net.parameters()) + + +def _stub_wrapper(backbone_lr_scale=None, frozen=False): + policy = nn.Module() + policy.encoders = nn.ModuleDict( + { + "front_img_1": _StubEncoder(frozen=frozen), + "other": nn.Linear(4, 4), # no backbone_parameters -> main group + } + ) + policy.trunk = nn.Linear(4, 4) + root = nn.Module() + root.nets = nn.ModuleDict({"policy": policy}) + + model_cfg = { + "optimizer": { + "_target_": "torch.optim.AdamW", + "_partial_": True, + "lr": 5e-5, + "weight_decay": 0.0001, + }, + "scheduler": None, + } + if backbone_lr_scale is not None: + model_cfg["backbone_lr_scale"] = backbone_lr_scale + config_tree = OmegaConf.create({"model": model_cfg}) + # A LightningModule cannot be built without a trainer / datamodule here, so + # call the real methods against the attributes they actually read. + stub = SimpleNamespace( + hparams=SimpleNamespace(config_tree=config_tree), + _as_config=ModelWrapper._as_config, + model=root, + trainer=SimpleNamespace(model=root), + ) + stub._backbone_param_groups = ModelWrapper._backbone_param_groups.__get__(stub) + return stub, root + + +@pytest.mark.parametrize("scale", [None, 1.0]) +def test_configure_optimizers_single_group_by_default(scale): + stub, _ = _stub_wrapper(backbone_lr_scale=scale) + out = ModelWrapper.configure_optimizers(stub) + groups = out["optimizer"].param_groups + assert len(groups) == 1 + assert groups[0]["lr"] == pytest.approx(5e-5) + + +def test_configure_optimizers_backbone_group(): + stub, root = _stub_wrapper(backbone_lr_scale=0.1) + out = ModelWrapper.configure_optimizers(stub) + groups = out["optimizer"].param_groups + assert len(groups) == 2 + by_lr = {round(g["lr"], 12): g for g in groups} + assert set(by_lr) == {round(5e-5, 12), round(5e-6, 12)} + + encoder = root.nets["policy"].encoders["front_img_1"] + backbone_ids = {id(p) for p in encoder.backbone_parameters()} + got_backbone = {id(p) for p in by_lr[round(5e-6, 12)]["params"]} + main = {id(p) for p in by_lr[round(5e-5, 12)]["params"]} + assert got_backbone == backbone_ids + assert not (got_backbone & main) + trainable = {id(p) for p in root.parameters() if p.requires_grad} + assert got_backbone | main == trainable + + +def test_configure_optimizers_skips_frozen_backbone_params(): + stub, root = _stub_wrapper(backbone_lr_scale=0.1, frozen=True) + out = ModelWrapper.configure_optimizers(stub) + every = [p for g in out["optimizer"].param_groups for p in g["params"]] + assert all(p.requires_grad for p in every) + assert {id(p) for p in every} == { + id(p) for p in root.parameters() if p.requires_grad + } + + +# -------------------------------------------------------------------------- +# Configs +# -------------------------------------------------------------------------- + +FLAGSHIP = "train_zarr_mecka_flagship_6d_hpt" + + +def test_base_model_config_declares_unit_backbone_lr_scale(compose_resolve): + cfg = compose_resolve(FLAGSHIP, []) + assert cfg.model.backbone_lr_scale == 1.0 + encoder = cfg.model.robomimic_model.encoder_specs.front_img_1 + assert encoder._target_ == "egomimic.models.hpt_nets.ResNet"