Method test_learner

axlearn/common/learner_test.py:1054–1333

Sets up two sub learners for encoder/decoder respectively.

(self, ema_decay: Optional[float], method: str)

    @parameterized.product(ema_decay=(None, 0.9), method=("update", "forward_and_backward"))
    # pylint: disable-next=too-many-statements
    def test_learner(self, ema_decay: Optional[float], method: str):
        """Sets up two sub learners for encoder/decoder respectively."""
        encoder_lr = 0.1
        opt1_cfg = config_for_function(sgd_optimizer).set(
            learning_rate=encoder_lr, decouple_weight_decay=True, weight_decay=1.0
        )
        opt2_cfg = config_for_function(adam_optimizer).set(
            learning_rate=0.0, b1=0.9, b2=0.99, eps=1e-5, l2_regularizer_weight=1.0
        )
        learner_rules = [(".*encoder.*", "encoder"), (".*decoder.*", "decoder")]

        cfg = CompositeLearner.default_config().set(
            name="test",
            rules=learner_rules,
            learners={
                "encoder": Learner.default_config().set(
                    optimizer=opt1_cfg, enable_per_variable_summaries=True
                ),
                "decoder": Learner.default_config().set(
                    optimizer=opt2_cfg, enable_per_variable_summaries=False
                ),
            },
        )
        cfg.ema.decay = ema_decay
        learner: CompositeLearner = cfg.instantiate(parent=None)

        param_specs = dict(
            encoder=dict(
                weight=ParameterSpec(
                    dtype=jnp.float32,
                    shape=[3, 4],
                    factorization=FactorizationSpec(axes=("row", "col")),
                    mesh_axes=PartitionSpec("data", "model"),
                ),
                bias=ParameterSpec(
                    dtype=jnp.float32,
                    shape=[4],
                    factorization=None,
                    mesh_axes=PartitionSpec("model"),
                ),
                mean=ParameterSpec(
                    dtype=jnp.float32,
                    shape=[2],
                    factorization=None,
                    mesh_axes=PartitionSpec("model"),
                ),
            ),
            decoder=dict(
                head=ParameterSpec(
                    dtype=jnp.float32,
                    shape=[2, 3],
                    factorization=FactorizationSpec(axes=("row", "col")),
                    mesh_axes=PartitionSpec("model", None),
                ),
                scalar=ParameterSpec(
                    dtype=jnp.float32,
                    shape=[],
                    factorization=None,