PaddleX/paddlex/modules/anomaly_detection/trainer.py

# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
#
# 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 os
import glob
from pathlib import Path
import lazy_paddle as paddle

from ..base import BaseTrainer, BaseTrainDeamon
from ...utils.config import AttrDict
from .model_list import MODELS


class UadTrainer(BaseTrainer):
    """Uad Model Trainer"""

    entities = MODELS

    def build_deamon(self, config: AttrDict) -> "SegTrainDeamon":
        """build deamon thread for saving training outputs timely

        Args:
            config (AttrDict): PaddleX pipeline config, which is loaded from pipeline yaml file.

        Returns:
            SegTrainDeamon: the training deamon thread object for saving training outputs timely.
        """
        return SegTrainDeamon(config)

    def update_config(self):
        """update training config"""
        self.pdx_config.update_dataset(self.global_config.dataset_dir, "SegDataset")
        if self.train_config.num_classes is not None:
            self.pdx_config.update_num_classes(self.train_config.num_classes)
        if (
            self.train_config.pretrain_weight_path
            and self.train_config.pretrain_weight_path != ""
        ):
            self.pdx_config.update_pretrained_weights(
                self.train_config.pretrain_weight_path, is_backbone=True
            )

    def get_train_kwargs(self) -> dict:
        """get key-value arguments of model training function

        Returns:
            dict: the arguments of training function.
        """
        train_args = {"device": self.get_device()}
        # XXX:
        os.environ.pop("FLAGS_npu_jit_compile", None)
        if self.train_config.batch_size is not None:
            train_args["batch_size"] = self.train_config.batch_size
        if self.train_config.learning_rate is not None:
            train_args["learning_rate"] = self.train_config.learning_rate
        if self.train_config.epochs_iters is not None:
            train_args["epochs_iters"] = self.train_config.epochs_iters
        if (
            self.train_config.resume_path is not None
            and self.train_config.resume_path != ""
        ):
            train_args["resume_path"] = self.train_config.resume_path
        if self.global_config.output is not None:
            train_args["save_dir"] = self.global_config.output
        if self.train_config.log_interval:
            train_args["log_iters"] = self.train_config.log_interval
        if self.train_config.eval_interval:
            train_args["do_eval"] = True
            train_args["save_interval"] = self.train_config.eval_interval
        train_args["dy2st"] = self.train_config.get("dy2st", False)
        return train_args


class SegTrainDeamon(BaseTrainDeamon):
    """SegTrainResultDemon"""

    last_k = 1

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def get_the_pdparams_suffix(self):
        """get the suffix of pdparams file"""
        return "pdparams"

    def get_the_pdema_suffix(self):
        """get the suffix of pdema file"""
        return "pdema"

    def get_the_pdopt_suffix(self):
        """get the suffix of pdopt file"""
        return "pdopt"

    def get_the_pdstates_suffix(self):
        """get the suffix of pdstates file"""
        return "pdstates"

    def get_ith_ckp_prefix(self, epoch_id):
        """get the prefix of the epoch_id checkpoint file"""
        return f"iter_{epoch_id}/model"

    def get_best_ckp_prefix(self):
        """get the prefix of the best checkpoint file"""
        return "best_model/model"

    def get_score(self, pdstates_path):
        """get the score by pdstates file"""
        if not Path(pdstates_path).exists():
            return 0
        return paddle.load(pdstates_path)["mIoU"]

    def get_epoch_id_by_pdparams_prefix(self, pdparams_dir):
        """get the epoch_id by pdparams file"""
        return int(pdparams_dir.parent.name.split("_")[-1])

    def update_result(self, result, train_output):
        """update every result"""
        train_output = Path(train_output).resolve()
        config_path = train_output.joinpath("config.yaml").resolve()
        if not config_path.exists():
            return result

        model_name = result["model_name"]
        if (
            model_name in self.config_recorder
            and self.config_recorder[model_name] != config_path
        ):
            result["models"] = self.init_model_pkg()
        result["config"] = config_path
        self.config_recorder[model_name] = config_path

        result["visualdl_log"] = self.update_vdl_log(train_output)
        result["label_dict"] = self.update_label_dict(train_output)

        model = self.get_model(result["model_name"], config_path)

        params_path_list = list(
            train_output.glob(
                ".".join(
                    [self.get_ith_ckp_prefix("[0-9]*"), self.get_the_pdparams_suffix()]
                )
            )
        )
        iter_ids = []
        for params_path in params_path_list:
            iter_id = self.get_epoch_id_by_pdparams_prefix(params_path)
            iter_ids.append(iter_id)
        iter_ids.sort()
        # TODO(gaotingquan): how to avoid that the latest ckp files is being saved
        # epoch_ids = epoch_ids[:-1]
        for i in range(1, self.last_k + 1):
            if len(iter_ids) < i:
                break
            self.update_models(
                result,
                model,
                train_output,
                f"last_{i}",
                self.get_ith_ckp_prefix(iter_ids[-i]),
            )
        self.update_models(
            result, model, train_output, "best", self.get_best_ckp_prefix()
        )
        return result

    def update_models(self, result, model, train_output, model_key, ckp_prefix):
        """update info of the models to be saved"""
        pdparams = train_output.joinpath(
            ".".join([ckp_prefix, self.get_the_pdparams_suffix()])
        )
        if pdparams.exists():
            recorder_key = f"{train_output.name}_{model_key}"
            if (
                model_key != "best"
                and recorder_key in self.model_recorder
                and self.model_recorder[recorder_key] == pdparams
            ):
                return

            self.model_recorder[recorder_key] = pdparams

            pdema = ""
            pdema_suffix = self.get_the_pdema_suffix()
            if pdema_suffix:
                pdema = pdparams.parents[1].joinpath(
                    ".".join([ckp_prefix, pdema_suffix])
                )
                if not pdema.exists():
                    pdema = ""

            pdopt = ""
            pdopt_suffix = self.get_the_pdopt_suffix()
            if pdopt_suffix:
                pdopt = pdparams.parents[1].joinpath(
                    ".".join([ckp_prefix, pdopt_suffix])
                )
                if not pdopt.exists():
                    pdopt = ""

            pdstates = ""
            pdstates_suffix = self.get_the_pdstates_suffix()
            if pdstates_suffix:
                pdstates = pdparams.parents[1].joinpath(
                    ".".join([ckp_prefix, pdstates_suffix])
                )
                if not pdstates.exists():
                    pdstates = ""

            score = self.get_score(Path(pdstates).resolve().as_posix())

            result["models"][model_key] = {
                "score": score,
                "pdparams": pdparams,
                "pdema": pdema,
                "pdopt": pdopt,
                "pdstates": pdstates,
            }

            self.update_inference_model(
                model,
                pdparams,
                train_output.joinpath(f"{ckp_prefix}"),
                result["models"][model_key],
            )