Method _learner_tree

axlearn/common/learner.py:434–452

Returns a tree of the same structure as params where each leaf is the name of the sublearner to apply.

(self, params: Nested[Any])

    def __init__(self, cfg: Config, *, parent: Module):
        super().__init__(cfg, parent=parent)
        cfg = self.config

        for name, learner_cfg in cfg.learners.items():
            # Sub learner should not hold ema.
            if learner_cfg.ema.decay is not None:
                raise ValueError(f"Sublearner {name} ema decay is not None.")
            if name == "ema":
                raise ValueError("Sublearner name cannot be ema.")

            sub_learner = learner_cfg.set(name=name)
            self._add_child(name, sub_learner)

        # Check that learners in the rules exist.
        for _, rule_name in cfg.rules:
            if rule_name not in cfg.learners:
                raise ValueError(f"{rule_name} is not found in the known learners.")
        if cfg.ema.decay is not None:
            # Create a global model ema.
            self.ema: PartitionedGradientTransformation = cfg.ema.instantiate()

    def _learner_tree(self, params: Nested[Any]) -> Nested[str]:
        """Returns a tree of the same structure as params where each leaf is the name of the