model.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299
  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 ...base import BaseModel
  13. from ...base.utils.arg import CLIArgument
  14. from ...base.utils.subprocess import CompletedProcess
  15. from ....utils.misc import abspath
  16. from ....utils import logging
  17. class ClsModel(BaseModel):
  18. """ Image Classification Model """
  19. def train(self,
  20. batch_size: int=None,
  21. learning_rate: float=None,
  22. epochs_iters: int=None,
  23. ips: str=None,
  24. device: str='gpu',
  25. resume_path: str=None,
  26. dy2st: bool=False,
  27. amp: str='OFF',
  28. num_workers: int=None,
  29. use_vdl: bool=True,
  30. save_dir: str=None,
  31. **kwargs) -> CompletedProcess:
  32. """train self
  33. Args:
  34. batch_size (int, optional): the train batch size value. Defaults to None.
  35. learning_rate (float, optional): the train learning rate value. Defaults to None.
  36. epochs_iters (int, optional): the train epochs value. Defaults to None.
  37. ips (str, optional): the ip addresses of nodes when using distribution. Defaults to None.
  38. device (str, optional): the running device. Defaults to 'gpu'.
  39. resume_path (str, optional): the checkpoint file path to resume training. Train from scratch if it is set
  40. to None. Defaults to None.
  41. dy2st (bool, optional): Enable dynamic to static. Defaults to False.
  42. amp (str, optional): the amp settings. Defaults to 'OFF'.
  43. num_workers (int, optional): the workers number. Defaults to None.
  44. use_vdl (bool, optional): enable VisualDL. Defaults to True.
  45. save_dir (str, optional): the directory path to save train output. Defaults to None.
  46. Returns:
  47. CompletedProcess: the result of training subprocess execution.
  48. """
  49. if resume_path is not None:
  50. resume_path = abspath(resume_path)
  51. with self._create_new_config_file() as config_path:
  52. # Update YAML config file
  53. config = self.config.copy()
  54. config._update_amp(amp)
  55. config._update_device(device)
  56. config._update_to_static(dy2st)
  57. config._update_use_vdl(use_vdl)
  58. if batch_size is not None:
  59. config.update_batch_size(batch_size)
  60. if learning_rate is not None:
  61. config.update_learning_rate(learning_rate)
  62. if epochs_iters is not None:
  63. config._update_epochs(epochs_iters)
  64. config._update_checkpoints(resume_path)
  65. if save_dir is not None:
  66. save_dir = abspath(save_dir)
  67. else:
  68. # `save_dir` is None
  69. save_dir = abspath(config.get_train_save_dir())
  70. config._update_output_dir(save_dir)
  71. if num_workers is not None:
  72. config.update_num_workers(num_workers)
  73. config.dump(config_path)
  74. cli_args = []
  75. do_eval = kwargs.pop('do_eval', True)
  76. profile = kwargs.pop('profile', None)
  77. if profile is not None:
  78. cli_args.append(CLIArgument('--profiler_options', profile))
  79. self._assert_empty_kwargs(kwargs)
  80. return self.runner.train(
  81. config_path, cli_args, device, ips, save_dir, do_eval=do_eval)
  82. def evaluate(self,
  83. weight_path: str,
  84. batch_size: int=None,
  85. ips: str=None,
  86. device: str='gpu',
  87. amp: str='OFF',
  88. num_workers: int=None,
  89. **kwargs) -> CompletedProcess:
  90. """evaluate self using specified weight
  91. Args:
  92. weight_path (str): the path of model weight file to be evaluated.
  93. batch_size (int, optional): the batch size value in evaluating. Defaults to None.
  94. ips (str, optional): the ip addresses of nodes when using distribution. Defaults to None.
  95. device (str, optional): the running device. Defaults to 'gpu'.
  96. amp (str, optional): the AMP setting. Defaults to 'OFF'.
  97. num_workers (int, optional): the workers number in evaluating. Defaults to None.
  98. Returns:
  99. CompletedProcess: the result of evaluating subprocess execution.
  100. """
  101. weight_path = abspath(weight_path)
  102. with self._create_new_config_file() as config_path:
  103. # Update YAML config file
  104. config = self.config.copy()
  105. config._update_amp(amp)
  106. config._update_device(device)
  107. config.update_pretrained_weights(weight_path)
  108. if batch_size is not None:
  109. config.update_batch_size(batch_size)
  110. if num_workers is not None:
  111. config.update_num_workers(num_workers)
  112. config.dump(config_path)
  113. self._assert_empty_kwargs(kwargs)
  114. cp = self.runner.evaluate(config_path, [], device, ips)
  115. return cp
  116. def predict(self,
  117. weight_path: str,
  118. input_path: str,
  119. input_list_path: str=None,
  120. device: str='gpu',
  121. save_dir: str=None,
  122. **kwargs) -> CompletedProcess:
  123. """predict using specified weight
  124. Args:
  125. weight_path (str): the path of model weight file used to predict.
  126. input_path (str): the path of image file to be predicted.
  127. input_list_path (str, optional): the paths of images to be predicted if is not None. Defaults to None.
  128. device (str, optional): the running device. Defaults to 'gpu'.
  129. save_dir (str, optional): the directory path to save predict output. Defaults to None.
  130. Returns:
  131. CompletedProcess: the result of predicting subprocess execution.
  132. """
  133. weight_path = abspath(weight_path)
  134. input_path = abspath(input_path)
  135. if input_list_path:
  136. input_list_path = abspath(input_list_path)
  137. with self._create_new_config_file() as config_path:
  138. # Update YAML config file
  139. config = self.config.copy()
  140. config.update_pretrained_weights(weight_path)
  141. config._update_predict_img(input_path, input_list_path)
  142. config._update_device(device)
  143. config._update_save_predict_result(save_dir)
  144. config.dump(config_path)
  145. self._assert_empty_kwargs(kwargs)
  146. return self.runner.predict(config_path, [], device)
  147. def export(self, weight_path: str, save_dir: str,
  148. **kwargs) -> CompletedProcess:
  149. """export the dynamic model to static model
  150. Args:
  151. weight_path (str): the model weight file path that used to export.
  152. save_dir (str): the directory path to save export output.
  153. Returns:
  154. CompletedProcess: the result of exporting subprocess execution.
  155. """
  156. weight_path = abspath(weight_path)
  157. save_dir = abspath(save_dir)
  158. with self._create_new_config_file() as config_path:
  159. # Update YAML config file
  160. config = self.config.copy()
  161. config.update_pretrained_weights(weight_path)
  162. config._update_save_inference_dir(save_dir)
  163. config.dump(config_path)
  164. self._assert_empty_kwargs(kwargs)
  165. return self.runner.export(config_path, [], None, save_dir)
  166. def infer(self,
  167. model_dir: str,
  168. input_path: str,
  169. device: str='gpu',
  170. save_dir: str=None,
  171. dict_path: str=None,
  172. **kwargs) -> CompletedProcess:
  173. """predict image using infernece model
  174. Args:
  175. model_dir (str): the directory path of inference model files that would use to predict.
  176. input_path (str): the path of image that would be predict.
  177. device (str, optional): the running device. Defaults to 'gpu'.
  178. save_dir (str, optional): the directory path to save output. Defaults to None.
  179. dict_path (str, optional): the label dict file path. Defaults to None.
  180. Returns:
  181. CompletedProcess: the result of infering subprocess execution.
  182. """
  183. model_dir = abspath(model_dir)
  184. input_path = abspath(input_path)
  185. if save_dir is not None:
  186. logging.warning("`save_dir` will not be used.")
  187. config_path = os.path.join(model_dir, 'inference.yml')
  188. config = self.config.copy()
  189. config.load(config_path)
  190. config._update_inference_model_dir(model_dir)
  191. config._update_infer_img(input_path)
  192. config._update_infer_device(device)
  193. if dict_path is not None:
  194. dict_path = abspath(dict_path)
  195. config.update_label_dict_path(dict_path)
  196. if 'enable_mkldnn' in kwargs:
  197. config._update_enable_mkldnn(kwargs.pop('enable_mkldnn'))
  198. with self._create_new_config_file() as config_path:
  199. config.dump(config_path)
  200. self._assert_empty_kwargs(kwargs)
  201. return self.runner.infer(config_path, [], device)
  202. def compression(self,
  203. weight_path: str,
  204. batch_size: int=None,
  205. learning_rate: float=None,
  206. epochs_iters: int=None,
  207. device: str='gpu',
  208. use_vdl: bool=True,
  209. save_dir: str=None,
  210. **kwargs) -> CompletedProcess:
  211. """compression model
  212. Args:
  213. weight_path (str): the path to weight file of model.
  214. batch_size (int, optional): the batch size value of compression training. Defaults to None.
  215. learning_rate (float, optional): the learning rate value of compression training. Defaults to None.
  216. epochs_iters (int, optional): the epochs or iters of compression training. Defaults to None.
  217. device (str, optional): the device to run compression training. Defaults to 'gpu'.
  218. use_vdl (bool, optional): whether or not to use VisualDL. Defaults to True.
  219. save_dir (str, optional): the directory to save output. Defaults to None.
  220. Returns:
  221. CompletedProcess: the result of compression subprocess execution.
  222. """
  223. weight_path = abspath(weight_path)
  224. with self._create_new_config_file() as config_path:
  225. # Update YAML config file
  226. config = self.config.copy()
  227. config._update_amp(None)
  228. config._update_device(device)
  229. config._update_use_vdl(use_vdl)
  230. config._update_slim_config(self.model_info[
  231. 'auto_compression_config_path'])
  232. config.update_pretrained_weights(weight_path)
  233. if batch_size is not None:
  234. config.update_batch_size(batch_size)
  235. if learning_rate is not None:
  236. config.update_learning_rate(learning_rate)
  237. if epochs_iters is not None:
  238. config._update_epochs(epochs_iters)
  239. if save_dir is not None:
  240. save_dir = abspath(save_dir)
  241. else:
  242. # `save_dir` is None
  243. save_dir = abspath(config.get_train_save_dir())
  244. config._update_output_dir(save_dir)
  245. config.dump(config_path)
  246. export_cli_args = []
  247. export_cli_args.append(
  248. CLIArgument(
  249. '-o',
  250. f"Global.save_inference_dir={os.path.join(save_dir, 'export')}"
  251. ))
  252. self._assert_empty_kwargs(kwargs)
  253. return self.runner.compression(config_path, [], export_cli_args,
  254. device, save_dir)