Method _run

caffe2/python/optimizer.py:430–478

(self, net, param_init_net, param_info)

        )

    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")

        # update (fused) in fp32
        net.MomentumSGDUpdate(
            [grad_fp32, momentum_data, lr, param_fp32],
            [grad_fp32, momentum_data, param_fp32],
            momentum=self.momentum,
            nesterov=self.nesterov,
        )

        # Copy updated param back to fp16
        net.FloatToHalf(param_fp32, param)


class FP16SgdOptimizer(SgdOptimizer):