PaddleX/paddlex/repo_apis/PaddleSeg_api/base_seg_config.py

# !/usr/bin/env python3
# -*- coding: UTF-8 -*-
################################################################################
#
# Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
#
################################################################################
"""
Author: PaddlePaddle Authors
"""
from urllib.parse import urlparse

import yaml
from paddleseg.utils import NoAliasDumper
from paddleseg.cvlibs.config import parse_from_yaml, merge_config_dicts

from ..base import BaseConfig
from ...utils.misc import abspath


class BaseSegConfig(BaseConfig):
    """ BaseSegConfig """

    def update(self, dict_like_obj):
        """ update """
        dict_ = merge_config_dicts(dict_like_obj, self.dict)
        self.reset_from_dict(dict_)

    def load(self, config_path):
        """ load """
        dict_ = parse_from_yaml(config_path)
        if not isinstance(dict_, dict):
            raise TypeError
        self.reset_from_dict(dict_)

    def dump(self, config_path):
        """ dump """
        with open(config_path, 'w', encoding='utf-8') as f:
            yaml.dump(self.dict, f, Dumper=NoAliasDumper)

    def update_learning_rate(self, learning_rate):
        """ update_learning_rate """
        if 'lr_scheduler' not in self:
            raise RuntimeError(
                "Not able to update learning rate, because no LR scheduler config was found."
            )
        self.lr_scheduler['learning_rate'] = learning_rate

    def update_batch_size(self, batch_size, mode='train'):
        """ update_batch_size """
        if mode == 'train':
            self.set_val('batch_size', batch_size)
        else:
            raise ValueError(
                f"Setting `batch_size` in {repr(mode)} mode is not supported.")

    def update_pretrained_weights(self, weight_path, is_backbone=False):
        """ update_pretrained_weights """
        if 'model' not in self:
            raise RuntimeError(
                "Not able to update pretrained weight path, because no model config was found."
            )
        if isinstance(weight_path, str):
            if urlparse(weight_path).scheme == '':
                # If `weight_path` is a string but not URL (with scheme present),
                # it will be recognized as a local file path.
                weight_path = abspath(weight_path)
        else:
            if weight_path is not None:
                raise TypeError("`weight_path` must be string or None.")
        if is_backbone:
            if 'backbone' not in self.model:
                raise RuntimeError(
                    "Not able to update pretrained weight path of backbone, because no backbone config was found."
                )
            self.model['backbone']['pretrained'] = weight_path
        else:
            self.model['pretrained'] = weight_path

    def update_dy2st(self, dy2st):
        """ update_dy2st """
        self.set_val('to_static_training', dy2st)

    def update_dataset(self, dataset_dir, dataset_type=None):
        """ update_dataset """
        raise NotImplementedError

    def get_epochs_iters(self):
        """ get_epochs_iters """
        raise NotImplementedError

    def get_learning_rate(self):
        """ get_learning_rate """
        raise NotImplementedError

    def get_batch_size(self, mode='train'):
        """ get_batch_size """
        raise NotImplementedError

    def get_qat_epochs_iters(self):
        """ get_qat_epochs_iters """
        return self.get_epochs_iters() // 2

    def get_qat_learning_rate(self):
        """ get_qat_learning_rate """
        return self.get_learning_rate() / 2