【PPSCI Doc No.12、13、14、15、16、17】ppsci.arch.Arch (#752)

* [Add] arch examples * [Change] examples * [Change] examples * [Change] register_input_transform * [Change] data with rand --------- Co-authored-by: HydrogenSulfate <490868991@qq.com>

Author: megemini

ppsci/arch/base.py

@@ -72,6 +72,23 @@ def concat_to_tensor(
 
  Returns:
  Tuple[paddle.Tensor, ...]: Concatenated tensor.
+
+ Examples:
+ >>> import paddle
+ >>> import ppsci
+ >>> model = ppsci.arch.Arch()
+ >>> # fetch one tensor
+ >>> out = model.concat_to_tensor({'x':paddle.rand([64, 64, 1])}, ('x',))
+ >>> print(out.dtype, out.shape)
+ paddle.float32 [64, 64, 1]
+ >>> # fetch more tensors
+ >>> out = model.concat_to_tensor(
+ ... {'x1':paddle.rand([64, 64, 1]), 'x2':paddle.rand([64, 64, 1])},
+ ... ('x1', 'x2'),
+ ... axis=2)
+ >>> print(out.dtype, out.shape)
+ paddle.float32 [64, 64, 2]
+
  """
  if len(keys) == 1:
  return data_dict[keys[0]]
@@ -90,6 +107,23 @@ def split_to_dict(
 
  Returns:
  Dict[str, paddle.Tensor]: Dict contains tensor.
+
+ Examples:
+ >>> import paddle
+ >>> import ppsci
+ >>> model = ppsci.arch.Arch()
+ >>> # split one tensor
+ >>> out = model.split_to_dict(paddle.rand([64, 64, 1]), ('x',))
+ >>> for k, v in out.items():
+ ... print(f"{k} {v.dtype} {v.shape}")
+ x paddle.float32 [64, 64, 1]
+ >>> # split more tensors
+ >>> out = model.split_to_dict(paddle.rand([64, 64, 2]), ('x1', 'x2'), axis=2)
+ >>> for k, v in out.items():
+ ... print(f"{k} {v.dtype} {v.shape}")
+ x1 paddle.float32 [64, 64, 1]
+ x2 paddle.float32 [64, 64, 1]
+
  """
  if len(keys) == 1:
  return {keys[0]: data_tensor}
@@ -105,6 +139,27 @@ def register_input_transform(
  Args:
  transform (Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]):
  Input transform of network, receive a single tensor dict and return a single tensor dict.
+
+ Examples:
+ >>> import ppsci
+ >>> def transform_in(in_):
+ ... x = in_["x"]
+ ... # transform input
+ ... x_ = 2.0 * x
+ ... input_trans = {"2x": x_}
+ ... return input_trans
+ >>> # `MLP` inherits from `Arch`
+ >>> model = ppsci.arch.MLP(
+ ... input_keys=("2x",),
+ ... output_keys=("y",),
+ ... num_layers=5,
+ ... hidden_size=32)
+ >>> model.register_input_transform(transform_in)
+ >>> out = model({"x":paddle.rand([64, 64, 1])})
+ >>> for k, v in out.items():
+ ... print(f"{k} {v.dtype} {v.shape}")
+ y paddle.float32 [64, 64, 1]
+
  """
  self._input_transform = transform
 
@@ -121,18 +176,61 @@ def register_output_transform(
  transform (Callable[[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]):
  Output transform of network, receive two single tensor dict(raw input
  and raw output) and return a single tensor dict(transformed output).
+
+ Examples:
+ >>> import ppsci
+ >>> def transform_out(in_, out):
+ ... x = in_["x"]
+ ... y = out["y"]
+ ... u = 2.0 * x * y
+ ... output_trans = {"u": u}
+ ... return output_trans
+ >>> # `MLP` inherits from `Arch`
+ >>> model = ppsci.arch.MLP(
+ ... input_keys=("x",),
+ ... output_keys=("y",),
+ ... num_layers=5,
+ ... hidden_size=32)
+ >>> model.register_output_transform(transform_out)
+ >>> out = model({"x":paddle.rand([64, 64, 1])})
+ >>> for k, v in out.items():
+ ... print(f"{k} {v.dtype} {v.shape}")
+ u paddle.float32 [64, 64, 1]
+
  """
  self._output_transform = transform
 
  def freeze(self):
- """Freeze all parameters."""
+ """Freeze all parameters.
+
+ Examples:
+ >>> import ppsci
+ >>> model = ppsci.arch.Arch()
+ >>> # freeze all parameters and make model `eval`
+ >>> model.freeze()
+ >>> assert not model.training
+ >>> for p in model.parameters():
+ ... assert p.stop_gradient
+
+ """
  for param in self.parameters():
  param.stop_gradient = True
 
  self.eval()
 
  def unfreeze(self):
- """Unfreeze all parameters."""
+ """Unfreeze all parameters.
+
+ Examples:
+ >>> import ppsci
+ >>> model = ppsci.arch.Arch()
+ >>> # unfreeze all parameters and make model `train`
+ >>> model.unfreeze()
+ >>> assert model.training
+ >>> for p in model.parameters():
+ ... assert not p.stop_gradient
+
+ """
  for param in self.parameters():
  param.stop_gradient = False