# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import numpy as np from op_test import get_places import paddle from paddle import base call_forward_post_hook = False call_forward_pre_hook = False class SimpleNet(paddle.nn.Layer): def __init__( self, hidden_size, vocab_size, num_steps=20, init_scale=0.1, is_sparse=False, dtype='float32', ): super().__init__() self.hidden_size = hidden_size self.vocab_size = vocab_size self.init_scale = init_scale self.num_steps = num_steps paddle.set_default_dtype(dtype) self.embedding = paddle.nn.Embedding( vocab_size, hidden_size, sparse=is_sparse, weight_attr=base.ParamAttr( name='embedding_para', initializer=paddle.nn.initializer.Uniform( low=-init_scale, high=init_scale ), ), ) self.softmax_bias = self.create_parameter( attr=base.ParamAttr(), shape=[self.vocab_size], dtype=dtype, default_initializer=paddle.nn.initializer.Uniform( low=-self.init_scale, high=self.init_scale ), ) def forward(self, input, label): x_emb = self.embedding(input) projection = paddle.matmul( x_emb, paddle.transpose(self.embedding.weight, perm=[1, 0]) ) projection = paddle.add(projection, self.softmax_bias) projection = paddle.reshape(projection, shape=[-1, self.vocab_size]) loss = paddle.nn.functional.softmax_with_cross_entropy( logits=projection, label=label, soft_label=False ) loss = paddle.reshape(loss, shape=[-1, self.num_steps]) loss = paddle.mean(loss, axis=[0]) loss = paddle.sum(loss) return loss def forward_post_hook(layer, input, output): global call_forward_post_hook call_forward_post_hook = True def forward_pre_hook(layer, input): global call_forward_pre_hook call_forward_pre_hook = True def forward_post_hook1(layer, input, output): return output * 2 def forward_pre_hook1(layer, input): input_return = (input[0] * 2, input[1]) return input_return def forward_pre_hook_with_kwargs(layer, args, kwargs): kwargs['x'] = kwargs['x'] * 2 return (args, kwargs) def forward_post_hook_with_kwargs(layer, inputs, kwargs, outputs): outputs = outputs + kwargs["x"] return outputs class SimpleNetWithKWArgs(paddle.nn.Layer): def __init__(self): super().__init__() def forward(self, x, y): z = x + y return z class DummyContextManager: def __init__(self, inp): self.input = inp def __enter__(self, *args, **kwargs): self.input.append(2) def __exit__(self, *args, **kwargs): self.input.append(-1) class FailsNetInForward(paddle.nn.Layer): def __init__(self) -> None: super().__init__() def forward(self, x, fail: bool = True): if fail: raise RuntimeError("failing in forward") return x class Test_Forward_Hook(unittest.TestCase): # test forward_pre_hook and forward_post_hook that have return value def test_forward_hook_return_value(self): seed = 90 for place in get_places(): with base.dygraph.guard(place): paddle.seed(seed) base.set_flags({'FLAGS_sort_sum_gradient': True}) input_word = ( np.array( [0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 3, 4, 5, 6, 7, 8] ) .reshape(6, 3) .astype('int64') ) input_word1 = input_word * 2 input_word = input_word.reshape((-1, 3, 1)) input_word1 = input_word1.reshape((-1, 3, 1)) y_data = ( np.array( [1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9] ) .reshape(6, 3) .astype('int64') ) y_data = y_data.reshape((-1, 1)) input = paddle.to_tensor(input_word) input1 = paddle.to_tensor(input_word1) y = paddle.to_tensor(y_data) simplenet = SimpleNet( hidden_size=20, vocab_size=32, num_steps=3, init_scale=0.1, is_sparse=False, dtype="float32", ) # origin, don't register any hook outs_origin = simplenet(input, y) outs_origin1 = simplenet(input1, y) # register forward_pre_hook forward_pre_hook_handle1 = simplenet.register_forward_pre_hook( forward_pre_hook1 ) outs_pre_hook = simplenet(input, y) np.testing.assert_array_equal( outs_pre_hook.numpy(), outs_origin1.numpy() ) # remove forward_pre_hook forward_pre_hook_handle1.remove() outs_pre_hook = simplenet(input, y) np.testing.assert_array_equal( outs_pre_hook.numpy(), outs_origin.numpy() ) # register forward_posst_hook forward_post_hook_handle1 = ( simplenet.register_forward_post_hook(forward_post_hook1) ) outs_forward_hook = simplenet(input, y) np.testing.assert_array_equal( outs_forward_hook.numpy(), outs_origin.numpy() * 2 ) # remove forward_post_hook forward_post_hook_handle1.remove() outs_forward_hook = simplenet(input, y) np.testing.assert_array_equal( outs_forward_hook.numpy(), outs_origin.numpy() ) # test forward_pre_hook and forward_post_hook that don't have return value def test_forward_hook(self): seed = 90 for place in get_places(): with base.dygraph.guard(place): paddle.seed(seed) base.set_flags({'FLAGS_sort_sum_gradient': True}) global call_forward_post_hook global call_forward_pre_hook input_word = ( np.array( [0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 3, 4, 5, 6, 7, 8] ) .reshape(6, 3) .astype('int64') ) input_word = input_word.reshape((-1, 3, 1)) y_data = ( np.array( [1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9] ) .reshape(6, 3) .astype('int64') ) y_data = y_data.reshape((-1, 1)) input = paddle.to_tensor(input_word) y = paddle.to_tensor(y_data) simplenet = SimpleNet( hidden_size=20, vocab_size=32, num_steps=3, init_scale=0.1, is_sparse=False, dtype="float32", ) # origin, don't register any hook outs_origin = simplenet(input, y) self.assertFalse(call_forward_post_hook) self.assertFalse(call_forward_pre_hook) # register forward_post_hook and forward_pre_hook forward_post_hook_handle = simplenet.register_forward_post_hook( forward_post_hook ) forward_pre_hook_handle = simplenet.register_forward_pre_hook( forward_pre_hook ) outs_hook = simplenet(input, y) self.assertTrue(call_forward_post_hook) self.assertTrue(call_forward_pre_hook) outs_hook = simplenet(input, y) self.assertTrue(call_forward_post_hook) self.assertTrue(call_forward_pre_hook) # remove forward_post_hook forward_post_hook_handle.remove() call_forward_post_hook = False call_forward_pre_hook = False outs_remove_forward_hook = simplenet(input, y) self.assertFalse(call_forward_post_hook) self.assertTrue(call_forward_pre_hook) # remove forward_pre_hook forward_pre_hook_handle.remove() call_forward_post_hook = False call_forward_pre_hook = False outs_remove_hook = simplenet(input, y) self.assertFalse(call_forward_post_hook) self.assertFalse(call_forward_pre_hook) def test_always_called_forward_hooks(self): x = paddle.ones((10, 10)) stack = [] ctx = None def setup_context(): nonlocal ctx ctx = DummyContextManager(stack) def ctx_setup_hook(m, i): setup_context() ctx.__enter__() def ctx_setup_failure_hook(m, i): setup_context() ctx.__enter__() raise RuntimeError("failing in ctx setup") def ctx_shutdown_hook(m, i, o): ctx.__exit__() def ctx_shutdown_failure_hook(m, i, o): ctx.__exit__() raise RuntimeError("failing in ctx shutdown") def throw_hook(m, i, o): raise RuntimeError("failing in throw") net = FailsNetInForward() forward_pre_hook_handle = net.register_forward_pre_hook(ctx_setup_hook) forward_post_hook_handle = net.register_forward_post_hook( ctx_shutdown_hook, always_call=True ) self.assertTrue(len(net._forward_post_hooks_always_called) == 1) # make sure always_called forward hook runs when model.forward raises RuntimeError with self.assertRaisesRegex(RuntimeError, "failing in forward"): net(x=x) self.assertEqual(stack, [2, -1]) # make sure that always_called forward hook does not run twice if there is no error net(x, fail=False) self.assertEqual(stack, [2, -1, 2, -1]) # make sure always_called forward hook runs when forward pre hook raises RuntimeError forward_pre_hook_handle.remove() net.register_forward_pre_hook(ctx_setup_failure_hook) with self.assertRaisesRegex(RuntimeError, "failing in ctx setup"): net(x, fail=False) self.assertEqual(stack, [2, -1, 2, -1, 2, -1]) # make sure always_called hook runs when another always_called forward hook raises an error forward_post_hook_handle2 = net.register_forward_post_hook( throw_hook, prepend=True, always_call=True ) # error raised should not be error of the forced hook with self.assertRaisesRegex(RuntimeError, "failing in ctx setup"): net(x, fail=False) self.assertEqual(stack, [2, -1, 2, -1, 2, -1, 2, -1]) # make sure that always called forward hooks are properly removed forward_post_hook_handle.remove() forward_post_hook_handle2.remove() self.assertTrue(len(net._forward_post_hooks_always_called) == 0) # make sure that always called forward hook is not run twice if it fails while running forward_post_hook_handle3 = net.register_forward_post_hook( ctx_shutdown_failure_hook, always_call=True ) with self.assertRaisesRegex(RuntimeError, "failing in ctx setup"): net(x, fail=False) self.assertEqual(stack, [2, -1, 2, -1, 2, -1, 2, -1, 2, -1]) class TestHookWithKWArgs(unittest.TestCase): def test_kwargs_hook(self): x = paddle.randn((2, 3)) y = paddle.randn((2, 3)) # 1. test forward pre hook net = SimpleNetWithKWArgs() remove_handler = net.register_forward_pre_hook( forward_pre_hook_with_kwargs, with_kwargs=True ) out = net(x=x, y=y) np.testing.assert_allclose(out.numpy(), (x * 2 + y).numpy()) remove_handler.remove() out = net(x=x, y=y) np.testing.assert_allclose(out.numpy(), (x + y).numpy()) # 2. test forward pre and forward post hooks net = SimpleNetWithKWArgs() net.register_forward_post_hook( forward_post_hook_with_kwargs, with_kwargs=True ) net.register_forward_pre_hook( forward_pre_hook_with_kwargs, with_kwargs=True ) out = net(x=x, y=y) np.testing.assert_allclose( out.numpy(), (x * 4 + y).numpy(), rtol=1e-5, atol=1e-6 ) def test_forward_hook_alias_and_prepend(self): x = paddle.ones((2, 3)) y = paddle.ones((2, 3)) net = SimpleNetWithKWArgs() hook_calls = [] def first_pre_hook(layer, args): hook_calls.append("first_pre") return (args[0] + 1, args[1]) def second_pre_hook(layer, args): hook_calls.append("second_pre") return (args[0] * 2, args[1]) def first_post_hook(layer, args, output): hook_calls.append("first_post") return output + 1 def second_post_hook(layer, args, output): hook_calls.append("second_post") return output * 2 net.register_forward_pre_hook(second_pre_hook) net.register_forward_pre_hook(first_pre_hook, prepend=True) net.register_forward_hook(second_post_hook) net.register_forward_hook(first_post_hook, prepend=True) out = net(x, y) self.assertEqual( hook_calls, ["first_pre", "second_pre", "first_post", "second_post"], ) np.testing.assert_allclose(out.numpy(), np.full((2, 3), 12.0)) def test_forward_pre_hook_with_kwargs_return_error(self): x = paddle.randn((2, 3)) y = paddle.randn((2, 3)) net = SimpleNetWithKWArgs() def invalid_pre_hook(layer, args, kwargs): return args net.register_forward_pre_hook(invalid_pre_hook, with_kwargs=True) with self.assertRaisesRegex( RuntimeError, "forward pre-hook must return None" ): net(x=x, y=y) class TestBackwardHook(unittest.TestCase): def test_backward_hooks(self): for place in get_places(): with base.dygraph.guard(place): class ParamOnlyLayer(paddle.nn.Layer): def __init__(self): super().__init__() self.weight = self.create_parameter( shape=[1], dtype="float32", is_bias=False ) def forward(self, x): return x * self.weight layer = ParamOnlyLayer() hook_calls = [] def full_backward_pre_hook(layer, grad_output): hook_calls.append(("pre", grad_output[0].numpy().copy())) def full_backward_hook(layer, grad_input, grad_output): hook_calls.append( ( "full", len(grad_input), grad_output[0].numpy().copy(), ) ) layer.register_full_backward_pre_hook(full_backward_pre_hook) layer.register_full_backward_hook(full_backward_hook) x = paddle.to_tensor([2.0], stop_gradient=True) y = layer(x) y.backward() self.assertEqual( [call[0] for call in hook_calls], ["pre", "full"] ) np.testing.assert_allclose(hook_calls[0][1], [1.0]) self.assertEqual(hook_calls[1][1], 0) np.testing.assert_allclose(hook_calls[1][2], [1.0]) np.testing.assert_allclose(layer.weight.grad.numpy(), [2.0]) with self.assertRaisesRegex( NotImplementedError, "Please use register_full_backward_hook instead", ): layer.register_backward_hook(lambda *args, **kwargs: None) def test_backward_hook_with_forward_pre_hook(self): for place in get_places(): with base.dygraph.guard(place): class PreHookLayer(paddle.nn.Layer): def forward(self, x): return x * 3 layer = PreHookLayer() hook_calls = [] def scale_input(layer, inputs): return (inputs[0] * 2,) def full_backward_hook(layer, grad_input, grad_output): hook_calls.append( ( grad_input[0].numpy().copy(), grad_output[0].numpy().copy(), ) ) layer.register_forward_pre_hook(scale_input) layer.register_full_backward_hook(full_backward_hook) x = paddle.to_tensor([1.0], stop_gradient=False) y = layer(x) y.backward() self.assertEqual(len(hook_calls), 1) np.testing.assert_allclose(hook_calls[0][0], [3.0]) np.testing.assert_allclose(hook_calls[0][1], [1.0]) np.testing.assert_allclose(x.grad.numpy(), [6.0]) def test_backward_hook_prepend_and_return(self): for place in get_places(): with base.dygraph.guard(place): class ScaleLayer(paddle.nn.Layer): def forward(self, x): return x * 2 layer = ScaleLayer() hook_calls = [] def full_backward_pre_hook(layer, grad_output): hook_calls.append(("pre1", grad_output[0].numpy().copy())) return (grad_output[0] * 2,) def full_backward_pre_hook_first(layer, grad_output): hook_calls.append(("pre2", grad_output[0].numpy().copy())) def full_backward_hook(layer, grad_input, grad_output): hook_calls.append( ( "full1", grad_input[0].numpy().copy(), grad_output[0].numpy().copy(), ) ) return (grad_input[0] * 3,) def full_backward_hook_first(layer, grad_input, grad_output): hook_calls.append( ( "full2", grad_input[0].numpy().copy(), grad_output[0].numpy().copy(), ) ) return (grad_input[0] * 5,) layer.register_full_backward_pre_hook(full_backward_pre_hook) layer.register_full_backward_pre_hook( full_backward_pre_hook_first, prepend=True ) layer.register_full_backward_hook(full_backward_hook) layer.register_full_backward_hook( full_backward_hook_first, prepend=True ) x = paddle.to_tensor([1.0], stop_gradient=False) y = layer(x) y.backward() self.assertEqual( [call[0] for call in hook_calls], ["pre2", "pre1", "full2", "full1"], ) np.testing.assert_allclose(hook_calls[0][1], [1.0]) np.testing.assert_allclose(hook_calls[1][1], [1.0]) np.testing.assert_allclose(hook_calls[2][1], [4.0]) np.testing.assert_allclose(hook_calls[2][2], [2.0]) np.testing.assert_allclose(hook_calls[3][1], [20.0]) np.testing.assert_allclose(hook_calls[3][2], [2.0]) np.testing.assert_allclose(x.grad.numpy(), [60.0]) self.assertEqual(len(layer._get_backward_hooks()), 2) def test_linear_full_backward_hook_result(self): for place in get_places(): with base.dygraph.guard(place): linear = paddle.nn.Linear(128, 64, bias_attr=False) weight = ( np.arange(128 * 64).reshape(128, 64).astype("float32") / 1000 ) linear.weight.set_value(weight) hook_calls = [] def full_backward_pre_hook(layer, grad_output): hook_calls.append(("pre", grad_output[0].numpy().copy())) return (grad_output[0] * 2,) def full_backward_hook(layer, grad_input, grad_output): hook_calls.append( ( "full", grad_input[0].numpy().copy(), grad_output[0].numpy().copy(), ) ) return (grad_input[0] * 3,) linear.register_full_backward_pre_hook(full_backward_pre_hook) linear.register_full_backward_hook(full_backward_hook) np_x = ( np.arange(2 * 128).reshape(2, 128).astype("float32") / 100 ) x = paddle.to_tensor(np_x, stop_gradient=False) y = linear(x) y.sum().backward() expected_grad_output = np.ones([2, 64], dtype="float32") expected_hook_grad_output = expected_grad_output * 2 expected_grad_input = np.matmul( expected_hook_grad_output, weight.T ) self.assertEqual( [call[0] for call in hook_calls], ["pre", "full"] ) np.testing.assert_allclose( hook_calls[0][1], expected_grad_output ) np.testing.assert_allclose( hook_calls[1][1], expected_grad_input, rtol=1e-5 ) np.testing.assert_allclose( hook_calls[1][2], expected_hook_grad_output ) np.testing.assert_allclose( x.grad.numpy(), expected_grad_input * 3, rtol=1e-5 ) if __name__ == '__main__': unittest.main()