model.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. from ...base import BaseModel
  16. from ...base.utils.arg import CLIArgument
  17. from ...base.utils.subprocess import CompletedProcess
  18. from ....utils.device import parse_device
  19. from ....utils.misc import abspath
  20. from ....utils import logging
  21. class TextRecModel(BaseModel):
  22. """Text Recognition Model"""
  23. METRICS = [
  24. "acc",
  25. "norm_edit_dis",
  26. "Teacher_acc",
  27. "Teacher_norm_edit_dis",
  28. "precision",
  29. "recall",
  30. "hmean",
  31. ]
  32. def train(
  33. self,
  34. batch_size: int = None,
  35. learning_rate: float = None,
  36. epochs_iters: int = None,
  37. ips: str = None,
  38. device: str = "gpu",
  39. resume_path: str = None,
  40. dy2st: bool = False,
  41. amp: str = "OFF",
  42. num_workers: int = None,
  43. use_vdl: bool = True,
  44. save_dir: str = None,
  45. **kwargs,
  46. ) -> CompletedProcess:
  47. """train self
  48. Args:
  49. batch_size (int, optional): the train batch size value. Defaults to None.
  50. learning_rate (float, optional): the train learning rate value. Defaults to None.
  51. epochs_iters (int, optional): the train epochs value. Defaults to None.
  52. ips (str, optional): the ip addresses of nodes when using distribution. Defaults to None.
  53. device (str, optional): the running device. Defaults to 'gpu'.
  54. resume_path (str, optional): the checkpoint file path to resume training. Train from scratch if it is set
  55. to None. Defaults to None.
  56. dy2st (bool, optional): Enable dynamic to static. Defaults to False.
  57. amp (str, optional): the amp settings. Defaults to 'OFF'.
  58. num_workers (int, optional): the workers number. Defaults to None.
  59. use_vdl (bool, optional): enable VisualDL. Defaults to True.
  60. save_dir (str, optional): the directory path to save train output. Defaults to None.
  61. Returns:
  62. CompletedProcess: the result of training subprocess execution.
  63. """
  64. config = self.config.copy()
  65. if batch_size is not None:
  66. config.update_batch_size(batch_size)
  67. if learning_rate is not None:
  68. config.update_learning_rate(learning_rate)
  69. if epochs_iters is not None:
  70. config._update_epochs(epochs_iters)
  71. # No need to handle `ips`
  72. config.update_device(device)
  73. if resume_path is not None:
  74. resume_path = abspath(resume_path)
  75. config._update_checkpoints(resume_path)
  76. config._update_to_static(dy2st)
  77. config._update_amp(amp)
  78. if num_workers is not None:
  79. config.update_num_workers(num_workers, "train")
  80. config._update_use_vdl(use_vdl)
  81. if save_dir is not None:
  82. save_dir = abspath(save_dir)
  83. else:
  84. save_dir = abspath(config.get_train_save_dir())
  85. config._update_output_dir(save_dir)
  86. cli_args = []
  87. do_eval = kwargs.pop("do_eval", True)
  88. profile = kwargs.pop("profile", None)
  89. if profile is not None:
  90. cli_args.append(CLIArgument("--profiler_options", profile))
  91. # Benchmarking mode settings
  92. benchmark = kwargs.pop("benchmark", None)
  93. if benchmark is not None:
  94. envs = benchmark.get("env", None)
  95. seed = benchmark.get("seed", None)
  96. do_eval = benchmark.get("do_eval", False)
  97. num_workers = benchmark.get("num_workers", None)
  98. config.update_log_ranks(device)
  99. config._update_amp(benchmark.get("amp", None))
  100. config.update_shuffle(benchmark.get("shuffle", False))
  101. config.update_cal_metrics(benchmark.get("cal_metrics", True))
  102. config.update_shared_memory(benchmark.get("shared_memory", True))
  103. config.update_print_mem_info(benchmark.get("print_mem_info", True))
  104. if num_workers is not None:
  105. config.update_num_workers(num_workers)
  106. if seed is not None:
  107. config.update_seed(seed)
  108. if envs is not None:
  109. for env_name, env_value in envs.items():
  110. os.environ[env_name] = str(env_value)
  111. # PDX related settings
  112. device_type = device.split(":")[0]
  113. if device_type in ["npu", "xpu", "mlu"]:
  114. uniform_output_enabled = False
  115. else:
  116. uniform_output_enabled = True
  117. config.update({"Global.uniform_output_enabled": uniform_output_enabled})
  118. config.update({"Global.pdx_model_name": self.name})
  119. hpi_config_path = self.model_info.get("hpi_config_path", None)
  120. config.update({"Global.hpi_config_path": hpi_config_path})
  121. self._assert_empty_kwargs(kwargs)
  122. with self._create_new_config_file() as config_path:
  123. config.dump(config_path)
  124. return self.runner.train(
  125. config_path, cli_args, device, ips, save_dir, do_eval=do_eval
  126. )
  127. def evaluate(
  128. self,
  129. weight_path: str,
  130. batch_size: int = None,
  131. ips: str = None,
  132. device: str = "gpu",
  133. amp: str = "OFF",
  134. num_workers: int = None,
  135. **kwargs,
  136. ) -> CompletedProcess:
  137. """evaluate self using specified weight
  138. Args:
  139. weight_path (str): the path of model weight file to be evaluated.
  140. batch_size (int, optional): the batch size value in evaluating. Defaults to None.
  141. ips (str, optional): the ip addresses of nodes when using distribution. Defaults to None.
  142. device (str, optional): the running device. Defaults to 'gpu'.
  143. amp (str, optional): the AMP setting. Defaults to 'OFF'.
  144. num_workers (int, optional): the workers number in evaluating. Defaults to None.
  145. Returns:
  146. CompletedProcess: the result of evaluating subprocess execution.
  147. """
  148. config = self.config.copy()
  149. weight_path = abspath(weight_path)
  150. config._update_checkpoints(weight_path)
  151. if batch_size is not None:
  152. config.update_batch_size(batch_size)
  153. # No need to handle `ips`
  154. config.update_device(device)
  155. config._update_amp(amp)
  156. if num_workers is not None:
  157. config.update_num_workers(num_workers, "eval")
  158. self._assert_empty_kwargs(kwargs)
  159. with self._create_new_config_file() as config_path:
  160. config.dump(config_path)
  161. cp = self.runner.evaluate(config_path, [], device, ips)
  162. return cp
  163. def predict(
  164. self,
  165. weight_path: str,
  166. input_path: str,
  167. device: str = "gpu",
  168. save_dir: str = None,
  169. **kwargs,
  170. ) -> CompletedProcess:
  171. """predict using specified weight
  172. Args:
  173. weight_path (str): the path of model weight file used to predict.
  174. input_path (str): the path of image file to be predicted.
  175. device (str, optional): the running device. Defaults to 'gpu'.
  176. save_dir (str, optional): the directory path to save predict output. Defaults to None.
  177. Returns:
  178. CompletedProcess: the result of predicting subprocess execution.
  179. """
  180. config = self.config.copy()
  181. weight_path = abspath(weight_path)
  182. config.update_pretrained_weights(weight_path)
  183. input_path = abspath(input_path)
  184. config._update_infer_img(
  185. input_path, infer_list=kwargs.pop("input_list_path", None)
  186. )
  187. config.update_device(device)
  188. # TODO: Handle `device`
  189. logging.warning("`device` will not be used.")
  190. if save_dir is not None:
  191. save_dir = abspath(save_dir)
  192. else:
  193. save_dir = abspath(config.get_predict_save_dir())
  194. config._update_save_res_path(os.path.join(save_dir, "res.txt"))
  195. self._assert_empty_kwargs(kwargs)
  196. with self._create_new_config_file() as config_path:
  197. config.dump(config_path)
  198. return self.runner.predict(config_path, [], device)
  199. def export(self, weight_path: str, save_dir: str, **kwargs) -> CompletedProcess:
  200. """export the dynamic model to static model
  201. Args:
  202. weight_path (str): the model weight file path that used to export.
  203. save_dir (str): the directory path to save export output.
  204. Returns:
  205. CompletedProcess: the result of exporting subprocess execution.
  206. """
  207. config = self.config.copy()
  208. device = kwargs.pop("device", None)
  209. if device:
  210. config.update_device(device)
  211. if not weight_path.startswith("http"):
  212. weight_path = abspath(weight_path)
  213. config.update_pretrained_weights(weight_path)
  214. save_dir = abspath(save_dir)
  215. config._update_save_inference_dir(save_dir)
  216. class_path = kwargs.pop("class_path", None)
  217. if class_path is not None:
  218. config.update_class_path(class_path)
  219. # PDX related settings
  220. config.update({"Global.pdx_model_name": self.name})
  221. hpi_config_path = self.model_info.get("hpi_config_path", None)
  222. config.update({"Global.hpi_config_path": hpi_config_path})
  223. self._assert_empty_kwargs(kwargs)
  224. with self._create_new_config_file() as config_path:
  225. config.dump(config_path)
  226. return self.runner.export(config_path, [], None, save_dir)
  227. def infer(
  228. self,
  229. model_dir: str,
  230. input_path: str,
  231. device: str = "gpu",
  232. save_dir: str = None,
  233. **kwargs,
  234. ) -> CompletedProcess:
  235. """predict image using infernece model
  236. Args:
  237. model_dir (str): the directory path of inference model files that would use to predict.
  238. input_path (str): the path of image that would be predict.
  239. device (str, optional): the running device. Defaults to 'gpu'.
  240. save_dir (str, optional): the directory path to save output. Defaults to None.
  241. Returns:
  242. CompletedProcess: the result of infering subprocess execution.
  243. """
  244. config = self.config.copy()
  245. cli_args = []
  246. model_dir = abspath(model_dir)
  247. cli_args.append(CLIArgument("--rec_model_dir", model_dir))
  248. input_path = abspath(input_path)
  249. cli_args.append(CLIArgument("--image_dir", input_path))
  250. device_type, _ = parse_device(device)
  251. cli_args.append(CLIArgument("--use_gpu", str(device_type == "gpu")))
  252. if save_dir is not None:
  253. logging.warning("`save_dir` will not be used.")
  254. dict_path = kwargs.pop("dict_path", None)
  255. if dict_path is not None:
  256. dict_path = abspath(dict_path)
  257. else:
  258. dict_path = config.get_label_dict_path()
  259. cli_args.append(CLIArgument("--rec_char_dict_path", dict_path))
  260. model_type = config._get_model_type()
  261. cli_args.append(CLIArgument("--rec_algorithm", model_type))
  262. infer_shape = config._get_infer_shape()
  263. if infer_shape is not None:
  264. cli_args.append(CLIArgument("--rec_image_shape", infer_shape))
  265. self._assert_empty_kwargs(kwargs)
  266. with self._create_new_config_file() as config_path:
  267. config.dump(config_path)
  268. return self.runner.infer(config_path, cli_args, device)
  269. def compression(
  270. self,
  271. weight_path: str,
  272. batch_size: int = None,
  273. learning_rate: float = None,
  274. epochs_iters: int = None,
  275. device: str = "gpu",
  276. use_vdl: bool = True,
  277. save_dir: str = None,
  278. **kwargs,
  279. ) -> CompletedProcess:
  280. """compression model
  281. Args:
  282. weight_path (str): the path to weight file of model.
  283. batch_size (int, optional): the batch size value of compression training. Defaults to None.
  284. learning_rate (float, optional): the learning rate value of compression training. Defaults to None.
  285. epochs_iters (int, optional): the epochs or iters of compression training. Defaults to None.
  286. device (str, optional): the device to run compression training. Defaults to 'gpu'.
  287. use_vdl (bool, optional): whether or not to use VisualDL. Defaults to True.
  288. save_dir (str, optional): the directory to save output. Defaults to None.
  289. Returns:
  290. CompletedProcess: the result of compression subprocess execution.
  291. """
  292. config = self.config.copy()
  293. export_cli_args = []
  294. weight_path = abspath(weight_path)
  295. config.update_pretrained_weights(weight_path)
  296. if batch_size is not None:
  297. config.update_batch_size(batch_size)
  298. if learning_rate is not None:
  299. config.update_learning_rate(learning_rate)
  300. if epochs_iters is not None:
  301. config._update_epochs(epochs_iters)
  302. config.update_device(device)
  303. config._update_use_vdl(use_vdl)
  304. if save_dir is not None:
  305. save_dir = abspath(save_dir)
  306. else:
  307. save_dir = abspath(config.get_train_save_dir())
  308. config._update_output_dir(save_dir)
  309. export_cli_args.append(
  310. CLIArgument(
  311. "-o", f"Global.save_inference_dir={os.path.join(save_dir, 'export')}"
  312. )
  313. )
  314. self._assert_empty_kwargs(kwargs)
  315. with self._create_new_config_file() as config_path:
  316. config.dump(config_path)
  317. return self.runner.compression(
  318. config_path, [], export_cli_args, device, save_dir
  319. )