Class MultiPrecisionSgdOptimizer

caffe2/python/optimizer.py:411–478

class MultiPrecisionSgdOptimizer(SgdOptimizer):
    def __init__(
        self,
        base_learning_rate=0.1,
        momentum=0.0,
        policy="fixed",
        nesterov=True,
        sparse_dedup_aggregator=None,
        **kwargs
    ):
        super().__init__(
            base_learning_rate=base_learning_rate,
            policy=policy,
            momentum=momentum,
            nesterov=nesterov,
            sparse_dedup_aggregator=sparse_dedup_aggregator,
            **kwargs
        )

    def _run(self, net, param_init_net, param_info):
        param = param_info.blob
        param_fp32 = (
            param_info.blob_copy[core.DataType.FLOAT]
            if param_info.blob_copy is not None
            else None
        )

        # If we have a straight fp32 parameter, run the base class
        if param_fp32 is None:
            return SgdOptimizer._run(self, net, param_init_net, param_info)

        grad = param_info.grad
        if self.base_learning_rate == 0:
            return
        assert (
            self.base_learning_rate > 0
        ), "Expect positive base learning rate, got {}".format(self.base_learning_rate)

        lr, _ = self.build_lr(
            net,
            param_init_net,
            base_learning_rate=-self.base_learning_rate,
            policy=self.policy,
            **(self.init_kwargs)
        )

        momentum_data = param_init_net.ConstantFill(
            param_fp32, str(param) + "_momentum", value=0.0
        )
        self._aux_params.local.append(momentum_data)

        assert not isinstance(
            grad, core.GradientSlice
        ), "MultiPrecisionSgd does not support sparse gradients"

        # Copy gradient to fp32
        grad_fp32 = net.HalfToFloat(grad, grad + "_fp32")