跳转至

Utils.expression(方程计算) 模块

ppsci.utils.expression

ExpressionSolver

Bases: Layer

Expression computing helper, which compute named result according to corresponding function and related inputs.

Examples:

>>> import ppsci
>>> model = ppsci.arch.MLP(("x", "y"), ("u", "v"), 5, 128)
>>> expr_solver = ExpressionSolver()
Source code in ppsci/utils/expression.py
class ExpressionSolver(nn.Layer):
    """Expression computing helper, which compute named result according to corresponding
    function and related inputs.

    Examples:
        >>> import ppsci
        >>> model = ppsci.arch.MLP(("x", "y"), ("u", "v"), 5, 128)
        >>> expr_solver = ExpressionSolver()
    """

    nvtx_flag: bool  # only for nsight analysis

    def __init__(self):
        super().__init__()

    def forward(self, *args, **kwargs):
        raise NotImplementedError(
            "Use train_forward/eval_forward/visu_forward instead of forward."
        )

    @jit.to_static
    def train_forward(
        self,
        expr_dicts: Tuple[Dict[str, Callable], ...],
        input_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
        model: arch.Arch,
        constraint: Dict[str, "constraint.Constraint"],
        label_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
        weight_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
    ) -> Tuple[Dict[str, "paddle.Tensor"], Dict[str, float]]:
        """Forward computation for training, including model forward and equation
        forward.

        Args:
            expr_dicts (Tuple[Dict[str, Callable], ...]): Tuple of expression dicts.
            input_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of input dicts.
            model (arch.Arch): NN model.
            constraint (Dict[str, "constraint.Constraint"]): Constraint dict.
            label_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of label dicts.
            weight_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of weight dicts.

        Returns:
            Tuple[Dict[str, "paddle.Tensor"], Dict[str, float]]:
                all_losses: A loss dictionary containing the output terms of all constraints,
                constraint_losses: The loss values of all constraints.
        """
        losses_all: Dict[str, "paddle.Tensor"] = {}
        losses_constraint: Dict[str, float] = {}

        for i, cst_name in enumerate(constraint):
            cst_obj = constraint[cst_name]

            # model forward
            if self.nvtx_flag:  # only for nsight analysis
                core.nvprof_nvtx_push(f"Constraint {cst_name}")

            output_dict = model(input_dicts[i])

            # equation forward
            data_dict = {k: v for k, v in input_dicts[i].items()}
            data_dict.update(output_dict)
            for name, expr in expr_dicts[i].items():
                output_dict[name] = expr(data_dict)

            # put field 'area' into output_dict
            if "area" in input_dicts[i]:
                output_dict["area"] = input_dicts[i]["area"]

            # clear differentiation cache
            clear()

            # compute loss for each constraint according to its' own output, label and weight
            losses: Dict[str, "paddle.Tensor"] = cst_obj.loss(
                output_dict,
                label_dicts[i],
                weight_dicts[i],
            )
            # update losses into 'losses_all' and 'losses_constraint'
            # 'losses_all': Will be send to loss aggregator for further computing final loss(scalar)
            # 'losses_constraint': Will be used in logging
            losses_constraint[cst_name] = 0.0
            for key in losses:
                losses_constraint[cst_name] += losses[key].item()
                if key in losses_all:
                    losses_all[key] += losses[key]
                else:
                    losses_all[key] = losses[key]

            if self.nvtx_flag:  # only for nsight analysis
                core.nvprof_nvtx_pop()

        return losses_all, losses_constraint

    @jit.to_static
    def eval_forward(
        self,
        expr_dict: Dict[str, Callable],
        input_dict: Dict[str, "paddle.Tensor"],
        model: arch.Arch,
        validator: "validate.Validator",
        label_dict: Dict[str, "paddle.Tensor"],
        weight_dict: Dict[str, "paddle.Tensor"],
    ) -> Tuple[Dict[str, "paddle.Tensor"], Dict[str, "paddle.Tensor"]]:
        """Forward computation for evaluation, including model forward and equation
        forward.

        Args:
            expr_dict (Dict[str, Callable]): Expression dict.
            input_dict (Dict[str, paddle.Tensor]): Input dict.
            model (arch.Arch): NN model.
            validator (validate.Validator): Validator.
            label_dict (Dict[str, paddle.Tensor]): Label dict.
            weight_dict (Dict[str, paddle.Tensor]): Weight dict.

        Returns:
            Tuple[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]]: Result dict and loss for
                given validator.
        """
        # model forward
        output_dict = model(input_dict)

        # equation forward
        data_dict = {k: v for k, v in input_dict.items()}
        data_dict.update(output_dict)
        for name, expr in expr_dict.items():
            output_dict[name] = expr(data_dict)

        # put field 'area' into output_dict
        if "area" in input_dict:
            output_dict["area"] = input_dict["area"]

        # clear differentiation cache
        clear()

        # compute loss for each validator according to its' own output, label and weight
        validator_losses = validator.loss(
            output_dict,
            label_dict,
            weight_dict,
        )
        return output_dict, validator_losses

    def visu_forward(
        self,
        expr_dict: Optional[Dict[str, Callable]],
        input_dict: Dict[str, "paddle.Tensor"],
        model: arch.Arch,
    ) -> Dict[str, "paddle.Tensor"]:
        """Forward computation for visualization, including model forward and equation
        forward.

        Args:
            expr_dict (Optional[Dict[str, Callable]]): Expression dict.
            input_dict (Dict[str, paddle.Tensor]): Input dict.
            model (arch.Arch): NN model.

        Returns:
            Dict[str, paddle.Tensor]: Result dict for given expression dict.
        """
        # model forward
        output_dict = model(input_dict)

        if isinstance(expr_dict, dict):
            # equation forward
            data_dict = {k: v for k, v in input_dict.items()}
            data_dict.update(output_dict)
            for name, expr in expr_dict.items():
                output_dict[name] = expr(data_dict)

        # clear differentiation cache
        clear()

        return output_dict
eval_forward(expr_dict, input_dict, model, validator, label_dict, weight_dict)

Forward computation for evaluation, including model forward and equation forward.

Parameters:

Name Type Description Default
expr_dict Dict[str, Callable]

Expression dict.

required
input_dict Dict[str, Tensor]

Input dict.

required
model Arch

NN model.

required
validator Validator

Validator.

required
label_dict Dict[str, Tensor]

Label dict.

required
weight_dict Dict[str, Tensor]

Weight dict.

required

Returns:

Type Description
Tuple[Dict[str, 'paddle.Tensor'], Dict[str, 'paddle.Tensor']]

Tuple[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]]: Result dict and loss for given validator.

Source code in ppsci/utils/expression.py
@jit.to_static
def eval_forward(
    self,
    expr_dict: Dict[str, Callable],
    input_dict: Dict[str, "paddle.Tensor"],
    model: arch.Arch,
    validator: "validate.Validator",
    label_dict: Dict[str, "paddle.Tensor"],
    weight_dict: Dict[str, "paddle.Tensor"],
) -> Tuple[Dict[str, "paddle.Tensor"], Dict[str, "paddle.Tensor"]]:
    """Forward computation for evaluation, including model forward and equation
    forward.

    Args:
        expr_dict (Dict[str, Callable]): Expression dict.
        input_dict (Dict[str, paddle.Tensor]): Input dict.
        model (arch.Arch): NN model.
        validator (validate.Validator): Validator.
        label_dict (Dict[str, paddle.Tensor]): Label dict.
        weight_dict (Dict[str, paddle.Tensor]): Weight dict.

    Returns:
        Tuple[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]]: Result dict and loss for
            given validator.
    """
    # model forward
    output_dict = model(input_dict)

    # equation forward
    data_dict = {k: v for k, v in input_dict.items()}
    data_dict.update(output_dict)
    for name, expr in expr_dict.items():
        output_dict[name] = expr(data_dict)

    # put field 'area' into output_dict
    if "area" in input_dict:
        output_dict["area"] = input_dict["area"]

    # clear differentiation cache
    clear()

    # compute loss for each validator according to its' own output, label and weight
    validator_losses = validator.loss(
        output_dict,
        label_dict,
        weight_dict,
    )
    return output_dict, validator_losses
train_forward(expr_dicts, input_dicts, model, constraint, label_dicts, weight_dicts)

Forward computation for training, including model forward and equation forward.

Parameters:

Name Type Description Default
expr_dicts Tuple[Dict[str, Callable], ...]

Tuple of expression dicts.

required
input_dicts Tuple[Dict[str, Tensor], ...]

Tuple of input dicts.

required
model Arch

NN model.

required
constraint Dict[str, 'constraint.Constraint']

Constraint dict.

required
label_dicts Tuple[Dict[str, Tensor], ...]

Tuple of label dicts.

required
weight_dicts Tuple[Dict[str, Tensor], ...]

Tuple of weight dicts.

required

Returns:

Type Description
Tuple[Dict[str, 'paddle.Tensor'], Dict[str, float]]

Tuple[Dict[str, "paddle.Tensor"], Dict[str, float]]: all_losses: A loss dictionary containing the output terms of all constraints, constraint_losses: The loss values of all constraints.

Source code in ppsci/utils/expression.py
@jit.to_static
def train_forward(
    self,
    expr_dicts: Tuple[Dict[str, Callable], ...],
    input_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
    model: arch.Arch,
    constraint: Dict[str, "constraint.Constraint"],
    label_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
    weight_dicts: Tuple[Dict[str, "paddle.Tensor"], ...],
) -> Tuple[Dict[str, "paddle.Tensor"], Dict[str, float]]:
    """Forward computation for training, including model forward and equation
    forward.

    Args:
        expr_dicts (Tuple[Dict[str, Callable], ...]): Tuple of expression dicts.
        input_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of input dicts.
        model (arch.Arch): NN model.
        constraint (Dict[str, "constraint.Constraint"]): Constraint dict.
        label_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of label dicts.
        weight_dicts (Tuple[Dict[str, paddle.Tensor], ...]): Tuple of weight dicts.

    Returns:
        Tuple[Dict[str, "paddle.Tensor"], Dict[str, float]]:
            all_losses: A loss dictionary containing the output terms of all constraints,
            constraint_losses: The loss values of all constraints.
    """
    losses_all: Dict[str, "paddle.Tensor"] = {}
    losses_constraint: Dict[str, float] = {}

    for i, cst_name in enumerate(constraint):
        cst_obj = constraint[cst_name]

        # model forward
        if self.nvtx_flag:  # only for nsight analysis
            core.nvprof_nvtx_push(f"Constraint {cst_name}")

        output_dict = model(input_dicts[i])

        # equation forward
        data_dict = {k: v for k, v in input_dicts[i].items()}
        data_dict.update(output_dict)
        for name, expr in expr_dicts[i].items():
            output_dict[name] = expr(data_dict)

        # put field 'area' into output_dict
        if "area" in input_dicts[i]:
            output_dict["area"] = input_dicts[i]["area"]

        # clear differentiation cache
        clear()

        # compute loss for each constraint according to its' own output, label and weight
        losses: Dict[str, "paddle.Tensor"] = cst_obj.loss(
            output_dict,
            label_dicts[i],
            weight_dicts[i],
        )
        # update losses into 'losses_all' and 'losses_constraint'
        # 'losses_all': Will be send to loss aggregator for further computing final loss(scalar)
        # 'losses_constraint': Will be used in logging
        losses_constraint[cst_name] = 0.0
        for key in losses:
            losses_constraint[cst_name] += losses[key].item()
            if key in losses_all:
                losses_all[key] += losses[key]
            else:
                losses_all[key] = losses[key]

        if self.nvtx_flag:  # only for nsight analysis
            core.nvprof_nvtx_pop()

    return losses_all, losses_constraint
visu_forward(expr_dict, input_dict, model)

Forward computation for visualization, including model forward and equation forward.

Parameters:

Name Type Description Default
expr_dict Optional[Dict[str, Callable]]

Expression dict.

required
input_dict Dict[str, Tensor]

Input dict.

required
model Arch

NN model.

required

Returns:

Type Description
Dict[str, 'paddle.Tensor']

Dict[str, paddle.Tensor]: Result dict for given expression dict.

Source code in ppsci/utils/expression.py
def visu_forward(
    self,
    expr_dict: Optional[Dict[str, Callable]],
    input_dict: Dict[str, "paddle.Tensor"],
    model: arch.Arch,
) -> Dict[str, "paddle.Tensor"]:
    """Forward computation for visualization, including model forward and equation
    forward.

    Args:
        expr_dict (Optional[Dict[str, Callable]]): Expression dict.
        input_dict (Dict[str, paddle.Tensor]): Input dict.
        model (arch.Arch): NN model.

    Returns:
        Dict[str, paddle.Tensor]: Result dict for given expression dict.
    """
    # model forward
    output_dict = model(input_dict)

    if isinstance(expr_dict, dict):
        # equation forward
        data_dict = {k: v for k, v in input_dict.items()}
        data_dict.update(output_dict)
        for name, expr in expr_dict.items():
            output_dict[name] = expr(data_dict)

    # clear differentiation cache
    clear()

    return output_dict