TPU gradient clipping.

pytorch_lightning/trainer/training_tricks.py

@@ -2,6 +2,7 @@
 from abc import ABC, abstractmethod
 
 import torch
+import math
 
 from pytorch_lightning.callbacks import GradientAccumulationScheduler
 
@@ -19,9 +20,30 @@ def get_model(self):
         pass
 
     def clip_gradients(self):
+        # this code is a modification of torch.nn.utils.clip_grad_norm_
+        # with TPU support based on https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.md
         if self.gradient_clip_val > 0:
             model = self.get_model()
-            torch.nn.utils.clip_grad_norm_(model.parameters(), self.gradient_clip_val)
+            parameters = model.parameters()
+            max_norm = self.gradient_clip_val
+            norm_type = 2
+            if isinstance(parameters, torch.Tensor):
+                parameters = [parameters]
+            parameters = list(filter(lambda p: p.grad is not None, parameters))
+            max_norm = float(max_norm)
+            norm_type = float(norm_type)
+            if norm_type == math.inf:
+                total_norm = max(p.grad.data.abs().max() for p in parameters)
+            else:
+                device = parameters[0].device
+                total_norm = torch.zeros([], device=device if parameters else None)
+                for p in parameters:
+                    param_norm = p.grad.data.norm(norm_type) ** norm_type
+                total_norm.add_(param_norm)
+                total_norm = (total_norm ** (1. / norm_type))
+            clip_coef = torch.tensor(max_norm, device=device) / (total_norm + 1e-6)
+            for p in parameters:
+                p.grad.data.mul_(torch.where(clip_coef < 1, clip_coef, torch.tensor(1., device=device)))
 
     def print_nan_gradients(self):
         model = self.get_model()