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]])
