|
|
|
|
|
|
|
|
|
|
| from caffe2.python import core |
| from caffe2.proto import caffe2_pb2 |
| from caffe2.python.optimizer import get_param_device |
| from caffe2.python.modeling.net_modifier import NetModifier |
|
|
| import logging |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class GradientClipping(NetModifier): |
|
|
| L1_NORM = 'l1_norm' |
| L2_NORM = 'l2_norm' |
|
|
| BY_NORM = 'by_norm' |
| BY_VALUE = 'by_value' |
|
|
| GRAD_CLIP_METHODS = [BY_NORM, BY_VALUE] |
| CLIP_GRADIENT_NORM_TYPES = [L2_NORM, L1_NORM] |
|
|
| def __init__(self, grad_clip_method, clip_norm_type='l2_norm', |
| clip_threshold=0.1, use_parameter_norm=False, |
| compute_norm_ratio=False, clip_max=1, clip_min=-1, |
| blobs_to_include=None, blobs_to_exclude=None): |
| """ |
| Clips gradient to avoid gradient magnitude explosion or vanishing gradient. |
| |
| Args: |
| grad_clip_method: ways to clip the gradients |
| clip_norm_type: type of norm used in the necessary computation |
| clip_threshold: threshold used to determine whether to clip |
| use_parameter_norm: a boolean to indicate whether to incorporate |
| the norm of the parameter |
| compute_norm_ratio: a boolean to compute the ratio between gradient norm |
| and parameter norm explicitly for debugging purpose |
| clip_max: when clipping by_value, any value that is greater than |
| clip_max will be clipped to clip_max |
| clip_min: when clipping by_value, any value that is smaller than |
| clip_min will be clipped to clip_min |
| blobs_to_include: names of blobs whose gradient is to be clipped. If it is set |
| to none, all param 's gradient in grad_map will be clipped. |
| blobs_to_exclude: names of blobs whose gradient is not to be clipped. |
| """ |
|
|
| assert grad_clip_method in self.GRAD_CLIP_METHODS, ( |
| "This method of clipping, {}, has not been implemented.".format( |
| clip_norm_type)) |
| if clip_norm_type is not None: |
| assert clip_norm_type in self.CLIP_GRADIENT_NORM_TYPES, ( |
| "This method of clipping, {}, has not been implemented.".format( |
| clip_norm_type)) |
|
|
| self.grad_clip_method = grad_clip_method |
| self.clip_norm_type = clip_norm_type |
| self.clip_threshold = float(clip_threshold) |
| self.use_parameter_norm = use_parameter_norm |
| self.compute_norm_ratio = compute_norm_ratio |
| self.clip_max = float(clip_max) |
| self.clip_min = float(clip_min) |
| self.blobs_to_include = blobs_to_include |
| self.blobs_to_exclude = blobs_to_exclude |
|
|
| def modify_net(self, net, init_net=None, grad_map=None, blob_to_device=None, |
| modify_output_record=False): |
|
|
| assert grad_map is not None |
|
|
| CPU = core.DeviceOption(caffe2_pb2.CPU) |
|
|
| final_param_map = {} |
| if self.blobs_to_include is None: |
| final_param_map = grad_map |
| else: |
| for blob in self.blobs_to_include: |
| param = core.BlobReference(blob) |
| if not net.BlobIsDefined(param): |
| raise Exception('param {0} is not defined in net {1}'.format( |
| param, net.Name())) |
| final_param_map[param] = grad_map[param] |
|
|
| if self.blobs_to_exclude is not None: |
| for blob in self.blobs_to_exclude: |
| final_param_map.pop(blob, None) |
|
|
| for param, grad in final_param_map.items(): |
| |
| |
| if isinstance(grad, core.GradientSlice): |
| continue |
|
|
| device = get_param_device( |
| param, |
| grad_map[str(param)], |
| param_to_device=blob_to_device, |
| default_device=CPU, |
| ) |
|
|
| with core.DeviceScope(device): |
| if self.grad_clip_method == self.BY_NORM: |
| if self.clip_norm_type == self.L2_NORM: |
| p = 2 |
| elif self.clip_norm_type == self.L1_NORM: |
| p = 1 |
|
|
| grad_norm = net.LpNorm( |
| [grad], |
| net.NextScopedBlob(prefix=str(grad) + '_l{}_norm'.format(p)), |
| p=p, |
| ) |
|
|
| if p == 2: |
| grad_norm = net.Pow([grad_norm], exponent=0.5) |
|
|
| op_inputs = [grad, grad_norm] |
|
|
| if self.use_parameter_norm: |
| param_norm = net.LpNorm( |
| [param], |
| net.NextScopedBlob( |
| prefix=str(param) + '_l{}_norm'.format(p)), |
| p=p, |
| ) |
|
|
| if p == 2: |
| param_norm = net.Pow([param_norm], exponent=0.5) |
|
|
| op_inputs.append(param_norm) |
|
|
| if self.compute_norm_ratio: |
| net.Div( |
| [grad_norm, param_norm], |
| [net.NextScopedBlob( |
| prefix=str(param) + "_norm_ratio")] |
| ) |
|
|
| net.ClipTensorByScaling( |
| op_inputs, |
| [grad], |
| threshold=self.clip_threshold, |
| ) |
| elif self.grad_clip_method == self.BY_VALUE: |
| net.Clip( |
| [grad], |
| [grad], |
| max=self.clip_max, |
| min=self.clip_min, |
| ) |
|
|
