model.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591
  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. import abc
  16. import inspect
  17. import functools
  18. import contextlib
  19. import tempfile
  20. import hashlib
  21. import base64
  22. from datetime import datetime, timedelta
  23. from .config import Config
  24. from .register import (
  25. get_registered_model_info,
  26. build_runner_from_model_info,
  27. build_model_from_model_info,
  28. )
  29. from ...utils import flags
  30. from ...utils import logging
  31. from ...utils.errors import (
  32. UnsupportedAPIError,
  33. UnsupportedParamError,
  34. raise_unsupported_api_error,
  35. )
  36. from ...utils.misc import CachedProperty as cached_property
  37. from ...utils.cache import get_cache_dir
  38. __all__ = ["PaddleModel", "BaseModel"]
  39. def _create_model(model_name=None, config=None):
  40. """_create_model"""
  41. if model_name is None and config is None:
  42. raise ValueError("At least one of `model_name` and `config` must be not None.")
  43. elif model_name is not None and config is not None:
  44. if model_name != config.model_name:
  45. raise ValueError(
  46. "If both `model_name` and `config` are not None, `model_name` should be the same as \
  47. `config.model_name`."
  48. )
  49. elif model_name is None and config is not None:
  50. model_name = config.model_name
  51. try:
  52. model_info = get_registered_model_info(model_name)
  53. except KeyError as e:
  54. raise UnsupportedParamError(
  55. f"{repr(model_name)} is not a registered model name."
  56. ) from e
  57. return build_model_from_model_info(model_info=model_info, config=config)
  58. PaddleModel = _create_model
  59. class BaseModel(metaclass=abc.ABCMeta):
  60. """
  61. Abstract base class of Model.
  62. Model defines how Config and Runner interact with each other. In addition,
  63. Model provides users with multiple APIs to perform model training,
  64. prediction, etc.
  65. """
  66. _API_FULL_LIST = ("train", "evaluate", "predict", "export", "infer", "compression")
  67. _API_SUPPORTED_OPTS_KEY_PATTERN = "supported_{api_name}_opts"
  68. def __init__(self, model_name, config=None):
  69. """
  70. Initialize the instance.
  71. Args:
  72. model_name (str): A registered model name.
  73. config (base.config.BaseConfig|None): Config. Default: None.
  74. """
  75. super().__init__()
  76. self.name = model_name
  77. self.model_info = get_registered_model_info(model_name)
  78. # NOTE: We build runner instance here by extracting runner info from model info
  79. # so that we don't have to overwrite the `__init__` method of each child class.
  80. self.runner = build_runner_from_model_info(self.model_info)
  81. if config is None:
  82. logging.warning(
  83. "We strongly discourage leaving `config` unset or setting it to None. "
  84. "Please note that when `config` is None, default settings will be used for every unspecified \
  85. configuration item, "
  86. "which may lead to unexpected result. Please make sure that this is what you intend to do."
  87. )
  88. config = Config(model_name)
  89. self.config = config
  90. self._patch_apis()
  91. @abc.abstractmethod
  92. def train(
  93. self,
  94. batch_size=None,
  95. learning_rate=None,
  96. epochs_iters=None,
  97. ips=None,
  98. device="gpu",
  99. resume_path=None,
  100. dy2st=False,
  101. amp="OFF",
  102. num_workers=None,
  103. use_vdl=True,
  104. save_dir=None,
  105. **kwargs,
  106. ):
  107. """
  108. Train a model.
  109. Args:
  110. batch_size (int|None): Number of samples in each mini-batch. If
  111. multiple devices are used, this is the batch size on each device.
  112. If None, use a default setting. Default: None.
  113. learning_rate (float|None): Learning rate of model training. If
  114. None, use a default setting. Default: None.
  115. epochs_iters (int|None): Total epochs or iterations of model
  116. training. If None, use a default setting. Default: None.
  117. ips (str|None): If not None, enable multi-machine training mode.
  118. `ips` specifies Paddle cluster node ips, e.g.,
  119. '192.168.0.16,192.168.0.17'. Default: None.
  120. device (str): A string that describes the device(s) to use, e.g.,
  121. 'cpu', 'gpu', 'gpu:1,2'. Default: 'gpu'.
  122. resume_path (str|None): If not None, resume training from the model
  123. snapshot corresponding to the weight file `resume_path`. If
  124. None, use a default setting. Default: None.
  125. dy2st (bool): Whether to enable dynamic-to-static training.
  126. Default: False.
  127. amp (str): Optimization level to use in AMP training. Choices are
  128. ['O1', 'O2', 'OFF']. Default: 'OFF'.
  129. num_workers (int|None): Number of subprocesses to use for data
  130. loading. If None, use a default setting. Default: None.
  131. use_vdl (bool): Whether to enable VisualDL during training.
  132. Default: True.
  133. save_dir (str|None): Directory to store model snapshots and logs. If
  134. None, use a default setting. Default: None.
  135. Returns:
  136. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  137. """
  138. raise NotImplementedError
  139. @abc.abstractmethod
  140. def evaluate(
  141. self,
  142. weight_path,
  143. batch_size=None,
  144. ips=None,
  145. device="gpu",
  146. amp="OFF",
  147. num_workers=None,
  148. **kwargs,
  149. ):
  150. """
  151. Evaluate a model.
  152. Args:
  153. weight_path (str): Path of the weights to initialize the model.
  154. batch_size (int|None): Number of samples in each mini-batch. If
  155. multiple devices are used, this is the batch size on each device.
  156. If None, use a default setting. Default: None.
  157. ips (str|None): If not None, enable multi-machine evaluation mode.
  158. `ips` specifies Paddle cluster node ips, e.g.,
  159. '192.168.0.16,192.168.0.17'. Default: None.
  160. device (str): A string that describes the device(s) to use, e.g.,
  161. 'cpu', 'gpu', 'gpu:1,2'. Default: 'gpu'.
  162. amp (str): Optimization level to use in AMP training. Choices are
  163. ['O1', 'O2', 'OFF']. Default: 'OFF'.
  164. num_workers (int|None): Number of subprocesses to use for data
  165. loading. If None, use a default setting. Default: None.
  166. Returns:
  167. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  168. """
  169. raise NotImplementedError
  170. @abc.abstractmethod
  171. def predict(self, weight_path, input_path, device="gpu", save_dir=None, **kwargs):
  172. """
  173. Make prediction with a pre-trained model.
  174. Args:
  175. weight_path (str): Path of the weights to initialize the model.
  176. input_path (str): Path of the input file, e.g. an image.
  177. device (str): A string that describes the device to use, e.g.,
  178. 'cpu', 'gpu'. Default: 'gpu'.
  179. save_dir (str|None): Directory to store prediction results. If None,
  180. use a default setting. Default: None.
  181. Returns:
  182. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  183. """
  184. raise NotImplementedError
  185. @abc.abstractmethod
  186. def export(self, weight_path, save_dir, **kwargs):
  187. """
  188. Export a pre-trained model.
  189. Args:
  190. weight_path (str): Path of the weights to initialize the model.
  191. save_dir (str): Directory to store the exported model.
  192. Returns:
  193. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  194. """
  195. raise NotImplementedError
  196. @abc.abstractmethod
  197. def infer(self, model_dir, input_path, device="gpu", save_dir=None, **kwargs):
  198. """
  199. Make inference with an exported inference model.
  200. Args:
  201. model_dir (str): Path of the exported inference model.
  202. input_path (str): Path of the input file, e.g. an image.
  203. device (str): A string that describes the device(s) to use, e.g.,
  204. 'cpu', 'gpu'. Default: 'gpu'.
  205. save_dir (str|None): Directory to store inference results. If None,
  206. use a default setting. Default: None.
  207. Returns:
  208. paddlex.repo_apis.base.utils.subprocess.CompletedProcess
  209. """
  210. raise NotImplementedError
  211. @abc.abstractmethod
  212. def compression(
  213. self,
  214. weight_path,
  215. batch_size=None,
  216. learning_rate=None,
  217. epochs_iters=None,
  218. device="gpu",
  219. use_vdl=True,
  220. save_dir=None,
  221. **kwargs,
  222. ):
  223. """
  224. Perform quantization aware training (QAT) and export the quantized
  225. model.
  226. Args:
  227. weight_path (str): Path of the weights to initialize the model.
  228. batch_size (int|None): Number of samples in each mini-batch. If
  229. multiple devices are used, this is the batch size on each
  230. device. If None, use a default setting. Default: None.
  231. learning_rate (float|None): Learning rate of QAT. If None, use a
  232. default setting. Default: None.
  233. epochs_iters (int|None): Total epochs of iterations of model
  234. training. If None, use a default setting. Default: None.
  235. device (str): A string that describes the device(s) to use, e.g.,
  236. 'cpu', 'gpu'. Default: 'gpu'.
  237. use_vdl (bool): Whether to enable VisualDL during training.
  238. Default: True.
  239. save_dir (str|None): Directory to store the results. If None, use a
  240. default setting. Default: None.
  241. Returns:
  242. tuple[paddlex.repo_apis.base.utils.subprocess.CompletedProcess]
  243. """
  244. raise NotImplementedError
  245. @contextlib.contextmanager
  246. def _create_new_config_file(self):
  247. cls = self.__class__
  248. model_name = self.model_info["model_name"]
  249. tag = "_".join([cls.__name__.lower(), model_name])
  250. yaml_file_name = tag + ".yml"
  251. if not flags.DEBUG:
  252. with tempfile.TemporaryDirectory(dir=get_cache_dir()) as td:
  253. path = os.path.join(td, yaml_file_name)
  254. with open(path, "w", encoding="utf-8"):
  255. pass
  256. yield path
  257. else:
  258. path = os.path.join(get_cache_dir(), yaml_file_name)
  259. with open(path, "w", encoding="utf-8"):
  260. pass
  261. yield path
  262. @contextlib.contextmanager
  263. def _create_new_val_json_file(self):
  264. cls = self.__class__
  265. model_name = self.model_info["model_name"]
  266. tag = "_".join([cls.__name__.lower(), model_name])
  267. json_file_name = tag + "_test.json"
  268. if not flags.DEBUG:
  269. with tempfile.TemporaryDirectory(dir=get_cache_dir()) as td:
  270. path = os.path.join(td, json_file_name)
  271. with open(path, "w", encoding="utf-8"):
  272. pass
  273. yield path
  274. else:
  275. path = os.path.join(get_cache_dir(), json_file_name)
  276. with open(path, "w", encoding="utf-8"):
  277. pass
  278. yield path
  279. @cached_property
  280. def supported_apis(self):
  281. """supported apis"""
  282. return self.model_info.get("supported_apis", None)
  283. @cached_property
  284. def supported_train_opts(self):
  285. """supported train opts"""
  286. return self.model_info.get(
  287. self._API_SUPPORTED_OPTS_KEY_PATTERN.format(api_name="train"), None
  288. )
  289. @cached_property
  290. def supported_evaluate_opts(self):
  291. """supported evaluate opts"""
  292. return self.model_info.get(
  293. self._API_SUPPORTED_OPTS_KEY_PATTERN.format(api_name="evaluate"), None
  294. )
  295. @cached_property
  296. def supported_predict_opts(self):
  297. """supported predcit opts"""
  298. return self.model_info.get(
  299. self._API_SUPPORTED_OPTS_KEY_PATTERN.format(api_name="predict"), None
  300. )
  301. @cached_property
  302. def supported_infer_opts(self):
  303. """supported infer opts"""
  304. return self.model_info.get(
  305. self._API_SUPPORTED_OPTS_KEY_PATTERN.format(api_name="infer"), None
  306. )
  307. @cached_property
  308. def supported_compression_opts(self):
  309. """supported copression opts"""
  310. return self.model_info.get(
  311. self._API_SUPPORTED_OPTS_KEY_PATTERN.format(api_name="compression"), None
  312. )
  313. @cached_property
  314. def supported_dataset_types(self):
  315. """supported dataset types"""
  316. return self.model_info.get("supported_dataset_types", None)
  317. @staticmethod
  318. def _assert_empty_kwargs(kwargs):
  319. if len(kwargs) > 0:
  320. # For compatibility
  321. logging.warning(f"Unconsumed keyword arguments detected: {kwargs}.")
  322. # raise RuntimeError(
  323. # f"Unconsumed keyword arguments detected: {kwargs}.")
  324. def _patch_apis(self):
  325. def _make_unavailable(bnd_method):
  326. @functools.wraps(bnd_method)
  327. def _unavailable_api(*args, **kwargs):
  328. model_name = self.name
  329. api_name = bnd_method.__name__
  330. raise UnsupportedAPIError(
  331. f"{model_name} does not support `{api_name}`."
  332. )
  333. return _unavailable_api
  334. def _add_prechecks(bnd_method):
  335. @functools.wraps(bnd_method)
  336. def _api_with_prechecks(*args, **kwargs):
  337. sig = inspect.Signature.from_callable(bnd_method)
  338. bnd_args = sig.bind(*args, **kwargs)
  339. args_dict = bnd_args.arguments
  340. # Merge default values
  341. for p in sig.parameters.values():
  342. if p.name not in args_dict and p.default is not p.empty:
  343. args_dict[p.name] = p.default
  344. # Rely on nonlocal variable `checks`
  345. for check in checks:
  346. # We throw any unhandled exception
  347. check.check(args_dict)
  348. return bnd_method(*args, **kwargs)
  349. api_name = bnd_method.__name__
  350. checks = []
  351. # We hardcode the prechecks for each API here
  352. if api_name == "train":
  353. opts = self.supported_train_opts
  354. if opts is not None:
  355. if "device" in opts:
  356. checks.append(
  357. _CheckDevice(
  358. opts["device"], self.runner.parse_device, check_mc=True
  359. )
  360. )
  361. if "dy2st" in opts:
  362. checks.append(_CheckDy2St(opts["dy2st"]))
  363. if "amp" in opts:
  364. checks.append(_CheckAMP(opts["amp"]))
  365. elif api_name == "evaluate":
  366. opts = self.supported_evaluate_opts
  367. if opts is not None:
  368. if "device" in opts:
  369. checks.append(
  370. _CheckDevice(
  371. opts["device"], self.runner.parse_device, check_mc=True
  372. )
  373. )
  374. if "amp" in opts:
  375. checks.append(_CheckAMP(opts["amp"]))
  376. elif api_name == "predict":
  377. opts = self.supported_predict_opts
  378. if opts is not None:
  379. if "device" in opts:
  380. checks.append(
  381. _CheckDevice(
  382. opts["device"], self.runner.parse_device, check_mc=False
  383. )
  384. )
  385. elif api_name == "infer":
  386. opts = self.supported_infer_opts
  387. if opts is not None:
  388. if "device" in opts:
  389. checks.append(
  390. _CheckDevice(
  391. opts["device"], self.runner.parse_device, check_mc=False
  392. )
  393. )
  394. elif api_name == "compression":
  395. opts = self.supported_compression_opts
  396. if opts is not None:
  397. if "device" in opts:
  398. checks.append(
  399. _CheckDevice(
  400. opts["device"], self.runner.parse_device, check_mc=True
  401. )
  402. )
  403. else:
  404. return bnd_method
  405. return _api_with_prechecks
  406. supported_apis = self.supported_apis
  407. if supported_apis is not None:
  408. avail_api_set = set(self.supported_apis)
  409. else:
  410. avail_api_set = set(self._API_FULL_LIST)
  411. for api_name in self._API_FULL_LIST:
  412. api = getattr(self, api_name)
  413. if api_name not in avail_api_set:
  414. # We decorate real API implementation with `_make_unavailable`
  415. # so that an error is always raised when the API is called.
  416. decorated_api = _make_unavailable(api)
  417. # Monkey-patch
  418. setattr(self, api_name, decorated_api)
  419. else:
  420. if flags.CHECK_OPTS:
  421. # We decorate real API implementation with `_add_prechecks`
  422. # to perform sanity and validity checks before invoking the
  423. # internal API.
  424. decorated_api = _add_prechecks(api)
  425. setattr(self, api_name, decorated_api)
  426. class _CheckFailed(Exception):
  427. """_CheckFailed"""
  428. # Allow `_CheckFailed` class to be recognized using `hasattr(exc, 'check_failed_error')`
  429. check_failed_error = True
  430. def __init__(self, arg_name, arg_val, legal_vals):
  431. self.arg_name = arg_name
  432. self.arg_val = arg_val
  433. self.legal_vals = legal_vals
  434. def __str__(self):
  435. return f"`{self.arg_name}` is expected to be one of or conforms to {self.legal_vals}, but got {self.arg_val}"
  436. class _APICallArgsChecker(object):
  437. """_APICallArgsChecker"""
  438. def __init__(self, legal_vals):
  439. super().__init__()
  440. self.legal_vals = legal_vals
  441. def check(self, args):
  442. """check"""
  443. raise NotImplementedError
  444. class _CheckDevice(_APICallArgsChecker):
  445. """_CheckDevice"""
  446. def __init__(self, legal_vals, parse_device, check_mc=False):
  447. super().__init__(legal_vals)
  448. self.parse_device = parse_device
  449. self.check_mc = check_mc
  450. def check(self, args):
  451. """check"""
  452. assert "device" in args
  453. device = args["device"]
  454. if device is not None:
  455. device_type, dev_ids = self.parse_device(device)
  456. if not self.check_mc:
  457. if device_type not in self.legal_vals:
  458. raise _CheckFailed("device", device, self.legal_vals)
  459. else:
  460. # Currently we only check multi-device settings for GPUs
  461. if device_type != "gpu":
  462. if device_type not in self.legal_vals:
  463. raise _CheckFailed("device", device, self.legal_vals)
  464. else:
  465. n1c1_desc = f"{device_type}_n1c1"
  466. n1cx_desc = f"{device_type}_n1cx"
  467. nxcx_desc = f"{device_type}_nxcx"
  468. if len(dev_ids) <= 1:
  469. if (
  470. n1c1_desc not in self.legal_vals
  471. and n1cx_desc not in self.legal_vals
  472. and nxcx_desc not in self.legal_vals
  473. ):
  474. raise _CheckFailed("device", device, self.legal_vals)
  475. else:
  476. assert "ips" in args
  477. if args["ips"] is not None:
  478. # Multi-machine
  479. if nxcx_desc not in self.legal_vals:
  480. raise _CheckFailed("device", device, self.legal_vals)
  481. else:
  482. # Single-machine multi-device
  483. if (
  484. n1cx_desc not in self.legal_vals
  485. and nxcx_desc not in self.legal_vals
  486. ):
  487. raise _CheckFailed("device", device, self.legal_vals)
  488. else:
  489. # When `device` is None, we assume that a default device that the
  490. # current model supports will be used, so we simply do nothing.
  491. pass
  492. class _CheckDy2St(_APICallArgsChecker):
  493. """_CheckDy2St"""
  494. def check(self, args):
  495. """check"""
  496. assert "dy2st" in args
  497. dy2st = args["dy2st"]
  498. if isinstance(self.legal_vals, list):
  499. assert len(self.legal_vals) == 1
  500. support_dy2st = bool(self.legal_vals[0])
  501. else:
  502. support_dy2st = bool(self.legal_vals)
  503. if dy2st is not None:
  504. if dy2st and not support_dy2st:
  505. raise _CheckFailed("dy2st", dy2st, [support_dy2st])
  506. else:
  507. pass
  508. class _CheckAMP(_APICallArgsChecker):
  509. """_CheckAMP"""
  510. def check(self, args):
  511. """check"""
  512. assert "amp" in args
  513. amp = args["amp"]
  514. if amp is not None:
  515. if amp != "OFF" and amp not in self.legal_vals:
  516. raise _CheckFailed("amp", amp, self.legal_vals)
  517. else:
  518. pass