| 15 | # if auto_clipper provided, compute max_norm using auto_clipper |
| 16 | # else, using give max_norm |
| 17 | def clip_grad_norm_(parameters, max_norm=1000000, norm_type=2, auto_clipper=None, loss_scale=1.0): |
| 18 | if isinstance(parameters, torch.Tensor): |
| 19 | parameters = [parameters] |
| 20 | parameters = list(filter(lambda p: p[1].grad is not None, parameters)) |
| 21 | |
| 22 | if len(parameters) == 0: return None |
| 23 | |
| 24 | max_norm = float(max_norm) |
| 25 | norm_type = float(norm_type) |
| 26 | if norm_type == inf: |
| 27 | total_norm = max(p.grad.data.abs().max() for p in parameters) |
| 28 | else: |
| 29 | total_norm = 0 |
| 30 | for name,p in parameters: |
| 31 | param_norm = p.grad.data.norm(norm_type) |
| 32 | total_norm += param_norm.item() ** norm_type |
| 33 | |
| 34 | total_norm = total_norm ** (1. / norm_type) |
| 35 | |
| 36 | # check inf/nan |
| 37 | overflow_num = torch.zeros(1) |
| 38 | if np.isinf(total_norm) or np.isnan(total_norm): |
| 39 | overflow_num[0] = 1 |
| 40 | dist.allreduce(overflow_num) |
| 41 | |
| 42 | if overflow_num > 0: |
| 43 | for name,p in parameters: |
| 44 | p.grad.data.fill_(float('nan')) |
| 45 | sync_print('total_norm is inf({})/nan({}), skip clipping!!!'.format(np.isinf(total_norm), np.isnan(total_norm))) |
| 46 | return total_norm |
| 47 | |
| 48 | # rescale the total_norm by loss_scale |
| 49 | total_norm /= loss_scale |
| 50 | |
| 51 | # update auto_clipper, compute max_norm |
| 52 | if auto_clipper is not None: |
| 53 | max_norm = auto_clipper.update(total_norm) |
| 54 | |
| 55 | # do clipping |
| 56 | clip_coef = max_norm / (total_norm + 1e-6) |
| 57 | if clip_coef < 1: |
| 58 | # sync_print('clip_coef: {}'.format(clip_coef)) |
| 59 | for _, p in parameters: |
| 60 | p.grad.data.mul_(clip_coef) |
| 61 | |
| 62 | return total_norm |
| 63 | |
| 64 | class ClipMeter(object): |
| 65 | def __init__(self, mom=None, thresh=None, min_max=False, mean=False, init=False): |