def enable_gradient_clipping(model: Model, grad_clipping: Optional[float]) -> None: if grad_clipping is not None: for parameter in model.parameters(): if parameter.requires_grad: parameter.register_hook(lambda grad: nn_util.clamp_tensor(grad, minimum=-grad_clipping, maximum=grad_clipping))
def enable_gradient_clipping(self) -> None: if self._grad_clipping is not None: for parameter in self.model.parameters(): if parameter.requires_grad: parameter.register_hook(lambda grad: nn_util.clamp_tensor( grad, minimum=-self._grad_clipping, maximum=self._grad_clipping))
def test_clamp_tensor(self): # Test on uncoalesced sparse tensor i = torch.LongTensor([[0, 1, 1, 0], [2, 0, 2, 2]]) v = torch.FloatTensor([3, 4, -5, 3]) tensor = torch.sparse.FloatTensor(i, v, torch.Size([2, 3])) clamped_tensor = util.clamp_tensor(tensor, minimum=-3, maximum=3).to_dense() assert_almost_equal(clamped_tensor, [[0, 0, 3], [3, 0, -3]]) # Test on coalesced sparse tensor i = torch.LongTensor([[0, 1, 1], [2, 0, 2]]) v = torch.FloatTensor([3, 4, -5]) tensor = torch.sparse.FloatTensor(i, v, torch.Size([2, 3])) clamped_tensor = util.clamp_tensor(tensor, minimum=-3, maximum=3).to_dense() assert_almost_equal(clamped_tensor, [[0, 0, 3], [3, 0, -3]]) # Test on dense tensor tensor = torch.tensor([[5, -4, 3], [-3, 0, -30]]) clamped_tensor = util.clamp_tensor(tensor, minimum=-3, maximum=3) assert_almost_equal(clamped_tensor, [[3, -3, 3], [-3, 0, -3]])