build_model.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. from ...repo_apis.base import Config, PaddleModel
  13. def build_model(model_name: str, config_path: str=None) -> tuple:
  14. """build Config and PaddleModel
  15. Args:
  16. model_name (str): model name
  17. config_path (str, optional): path to the PaddleX config yaml file.
  18. Defaults to None, i.e. using the default config file.
  19. Returns:
  20. tuple(Config, PaddleModel): the Config and PaddleModel
  21. """
  22. config = Config(model_name, config_path)
  23. # NOTE(gaotingquan): aistudio can only support 4GB shm when single gpu can be used.
  24. device = os.environ.get("DEVICE", None)
  25. if device:
  26. if device.lower() == "cpu" or (
  27. device.lower() == "gpu" and
  28. os.environ.get("GPU_NUMBER", None) == "1"):
  29. if hasattr(config, "disable_shared_memory"):
  30. config.disable_shared_memory()
  31. if hasattr(config, "update_num_workers"):
  32. config.update_num_workers(2)
  33. model = PaddleModel(config=config)
  34. return config, model