Class FP16SgdOptimizer

caffe2/python/optimizer.py:481–586

class FP16SgdOptimizer(SgdOptimizer):
    def __init__(
        self,
        base_learning_rate=0.1,
        momentum=0.0,
        policy="fixed",
        nesterov=True,
        weight_decay=0.0001,
        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
        )
        self.weight_decay = weight_decay

    def _run(self, net, param_init_net, param_info, fp32_update=False):

        fp32_update_flag = 0
        param_name = str(param_info.blob)

        # should only be triggered in FP16 training by SpatialBN, which
        # requires FP32 params in CuDNN.
        if param_name.find("spatbn") != -1:
            fp32_update = True

        if fp32_update:
            # doing a 32bit update
            # Have to assume param_info.blob is FP32 as there is no way
            # (that i currently know of) to query a blob's type in python
            fp32_update_flag = 1
            param = param_info.blob
            param_fp32 = param_info.blob
        else:
            if param_info.blob_copy is None:
                # doing a 32bit update
                # Have to assume param_info.blob is FP32 as there is no way
                # (that i currently know of) to query a blob's type in python
                fp32_update_flag = 1
                param = param_info.blob
                param_fp32 = param_info.blob
            else:
                if core.DataType.FLOAT in param_info.blob_copy:
                    param = param_info.blob
                    param_fp32 = param_info.blob_copy[core.DataType.FLOAT]
                elif core.DataType.FLOAT16 in param_info.blob_copy:
                    param = param_info.blob_copy[core.DataType.FLOAT16]
                    param_fp32 = param_info.blob
                else:
                    AssertionError(
                        "Unrecognized parameter format to be updated "
                        "by FP16 Optimizer. Parameter: {}".format(param_info.name)
                    )