[Dy2stat]Allow users to switch eval/train mode when using @to_static to decorate a function (#37383)

* Allow users to switch eval/train mode when using @to_static to decorate a function * refine code for train() and eval()

Author: 0x45f

python/paddle/fluid/dygraph/dygraph_to_static/program_translator.py

@@ -261,6 +261,25 @@ def __init__(self, function, input_spec=None, **kwargs):
  # Note: Hold a reference to ProgramTranslator for switching `enable_to_static`.
  self._program_trans = ProgramTranslator()
  self._kwargs = kwargs
+ self._training = True
+
+ def train(self):
+ if isinstance(self._class_instance,
+ layers.Layer) and self._class_instance.training == False:
+ raise RuntimeError(
+ "Failed to switch train mode. {} is a Layer's method, "
+ "please use Layer.train() to switch train mode.".format(
+ self.dygraph_function))
+ self._training = True
+
+ def eval(self):
+ if isinstance(self._class_instance,
+ layers.Layer) and self._class_instance.training == True:
+ raise RuntimeError(
+ "Failed to switch eval mode. {} is a Layer's method, "
+ "please use Layer.eval() to switch eval mode.".format(
+ self.dygraph_function))
+ self._training = False
 
  def __get__(self, instance, owner):
  """
@@ -340,6 +359,8 @@ def __call__(self, *args, **kwargs):
  # 3. synchronize self.training attribute.
  if isinstance(self._class_instance, layers.Layer):
  partial_program_layer.training = self._class_instance.training
+ else:
+ partial_program_layer.training = self._training
 
  # 4. return outputs.
  try:

python/paddle/fluid/tests/unittests/dygraph_to_static/test_program_translator.py

@@ -297,5 +297,52 @@ def test_raise_error(self):
  self.program_translator.get_program(net.forward, self.x)
 
 
+class SwitchModeNet(paddle.nn.Layer):
+ def __init__(self):
+ super(SwitchModeNet, self).__init__()
+
+ @paddle.jit.to_static
+ def forward(self, x):
+ return x + 1
+
+ @paddle.jit.to_static
+ def foo(self):
+ return True
+
+
+@paddle.jit.to_static
+def switch_mode_funciton():
+ return True
+
+
+class TestFunctionTrainEvalMode(unittest.TestCase):
+ def test_switch_mode(self):
+ paddle.disable_static()
+ switch_mode_funciton.eval()
+ switch_mode_funciton()
+ self.assertEqual(switch_mode_funciton._training, False)
+ _, partial_layer = switch_mode_funciton.program_cache.last()[-1]
+ self.assertEqual(partial_layer.training, False)
+
+ switch_mode_funciton.train()
+ switch_mode_funciton()
+ self.assertEqual(switch_mode_funciton._training, True)
+ _, partial_layer = switch_mode_funciton.program_cache.last()[-1]
+ self.assertEqual(partial_layer.training, True)
+
+ def test_raise_error(self):
+ paddle.disable_static()
+ net = SwitchModeNet()
+
+ self.assertEqual(net.training, True)
+ with self.assertRaises(RuntimeError):
+ net.forward.eval()
+
+ net.eval()
+ self.assertEqual(net.training, False)
+ with self.assertRaises(RuntimeError):
+ net.foo.train()
+
+
 if __name__ == '__main__':
  unittest.main()