callbacks.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  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. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. import os
  18. import sys
  19. import datetime
  20. import six
  21. import copy
  22. import json
  23. import paddle
  24. import paddle.distributed as dist
  25. from paddlex.ppdet.utils.checkpoint import save_model
  26. from paddlex.ppdet.metrics import get_infer_results
  27. from paddlex.ppdet.utils.logger import setup_logger
  28. logger = setup_logger('ppdet.engine')
  29. __all__ = [
  30. 'Callback', 'ComposeCallback', 'LogPrinter', 'Checkpointer',
  31. 'VisualDLWriter', 'SniperProposalsGenerator'
  32. ]
  33. class Callback(object):
  34. def __init__(self, model):
  35. self.model = model
  36. def on_step_begin(self, status):
  37. pass
  38. def on_step_end(self, status):
  39. pass
  40. def on_epoch_begin(self, status):
  41. pass
  42. def on_epoch_end(self, status):
  43. pass
  44. def on_train_begin(self, status):
  45. pass
  46. def on_train_end(self, status):
  47. pass
  48. class ComposeCallback(object):
  49. def __init__(self, callbacks):
  50. callbacks = [c for c in list(callbacks) if c is not None]
  51. for c in callbacks:
  52. assert isinstance(
  53. c, Callback), "callback should be subclass of Callback"
  54. self._callbacks = callbacks
  55. def on_step_begin(self, status):
  56. for c in self._callbacks:
  57. c.on_step_begin(status)
  58. def on_step_end(self, status):
  59. for c in self._callbacks:
  60. c.on_step_end(status)
  61. def on_epoch_begin(self, status):
  62. for c in self._callbacks:
  63. c.on_epoch_begin(status)
  64. def on_epoch_end(self, status):
  65. for c in self._callbacks:
  66. c.on_epoch_end(status)
  67. def on_train_begin(self, status):
  68. for c in self._callbacks:
  69. c.on_train_begin(status)
  70. def on_train_end(self, status):
  71. for c in self._callbacks:
  72. c.on_train_end(status)
  73. class LogPrinter(Callback):
  74. def __init__(self, model):
  75. super(LogPrinter, self).__init__(model)
  76. def on_step_end(self, status):
  77. if dist.get_world_size() < 2 or dist.get_rank() == 0:
  78. mode = status['mode']
  79. if mode == 'train':
  80. epoch_id = status['epoch_id']
  81. step_id = status['step_id']
  82. steps_per_epoch = status['steps_per_epoch']
  83. training_staus = status['training_staus']
  84. batch_time = status['batch_time']
  85. data_time = status['data_time']
  86. epoches = self.model.cfg.epoch
  87. batch_size = self.model.cfg['{}Reader'.format(mode.capitalize(
  88. ))]['batch_size']
  89. logs = training_staus.log()
  90. space_fmt = ':' + str(len(str(steps_per_epoch))) + 'd'
  91. if step_id % self.model.cfg.log_iter == 0:
  92. eta_steps = (epoches - epoch_id
  93. ) * steps_per_epoch - step_id
  94. eta_sec = eta_steps * batch_time.global_avg
  95. eta_str = str(datetime.timedelta(seconds=int(eta_sec)))
  96. ips = float(batch_size) / batch_time.avg
  97. fmt = ' '.join([
  98. 'Epoch: [{}]',
  99. '[{' + space_fmt + '}/{}]',
  100. 'learning_rate: {lr:.6f}',
  101. '{meters}',
  102. 'eta: {eta}',
  103. 'batch_cost: {btime}',
  104. 'data_cost: {dtime}',
  105. 'ips: {ips:.4f} images/s',
  106. ])
  107. fmt = fmt.format(
  108. epoch_id,
  109. step_id,
  110. steps_per_epoch,
  111. lr=status['learning_rate'],
  112. meters=logs,
  113. eta=eta_str,
  114. btime=str(batch_time),
  115. dtime=str(data_time),
  116. ips=ips)
  117. logger.info(fmt)
  118. if mode == 'eval':
  119. step_id = status['step_id']
  120. if step_id % 100 == 0:
  121. logger.info("Eval iter: {}".format(step_id))
  122. def on_epoch_end(self, status):
  123. if dist.get_world_size() < 2 or dist.get_rank() == 0:
  124. mode = status['mode']
  125. if mode == 'eval':
  126. sample_num = status['sample_num']
  127. cost_time = status['cost_time']
  128. logger.info('Total sample number: {}, averge FPS: {}'.format(
  129. sample_num, sample_num / cost_time))
  130. class Checkpointer(Callback):
  131. def __init__(self, model):
  132. super(Checkpointer, self).__init__(model)
  133. cfg = self.model.cfg
  134. self.best_ap = 0.
  135. self.save_dir = os.path.join(self.model.cfg.save_dir,
  136. self.model.cfg.filename)
  137. if hasattr(self.model.model, 'student_model'):
  138. self.weight = self.model.model.student_model
  139. else:
  140. self.weight = self.model.model
  141. def on_epoch_end(self, status):
  142. # Checkpointer only performed during training
  143. mode = status['mode']
  144. epoch_id = status['epoch_id']
  145. weight = None
  146. save_name = None
  147. if dist.get_world_size() < 2 or dist.get_rank() == 0:
  148. if mode == 'train':
  149. end_epoch = self.model.cfg.epoch
  150. if (
  151. epoch_id + 1
  152. ) % self.model.cfg.snapshot_epoch == 0 or epoch_id == end_epoch - 1:
  153. save_name = str(
  154. epoch_id
  155. ) if epoch_id != end_epoch - 1 else "model_final"
  156. weight = self.weight.state_dict()
  157. elif mode == 'eval':
  158. if 'save_best_model' in status and status['save_best_model']:
  159. for metric in self.model._metrics:
  160. map_res = metric.get_results()
  161. if 'bbox' in map_res:
  162. key = 'bbox'
  163. elif 'keypoint' in map_res:
  164. key = 'keypoint'
  165. else:
  166. key = 'mask'
  167. if key not in map_res:
  168. logger.warning("Evaluation results empty, this may be due to " \
  169. "training iterations being too few or not " \
  170. "loading the correct weights.")
  171. return
  172. if map_res[key][0] >= self.best_ap:
  173. self.best_ap = map_res[key][0]
  174. save_name = 'best_model'
  175. weight = self.weight.state_dict()
  176. logger.info("Best test {} ap is {:0.3f}.".format(
  177. key, self.best_ap))
  178. if weight:
  179. if self.model.use_ema:
  180. # save model and ema_model
  181. save_model(
  182. status['weight'],
  183. self.model.optimizer,
  184. self.save_dir,
  185. save_name,
  186. epoch_id + 1,
  187. ema_model=weight)
  188. else:
  189. save_model(weight, self.model.optimizer, self.save_dir,
  190. save_name, epoch_id + 1)
  191. class WiferFaceEval(Callback):
  192. def __init__(self, model):
  193. super(WiferFaceEval, self).__init__(model)
  194. def on_epoch_begin(self, status):
  195. assert self.model.mode == 'eval', \
  196. "WiferFaceEval can only be set during evaluation"
  197. for metric in self.model._metrics:
  198. metric.update(self.model.model)
  199. sys.exit()
  200. class VisualDLWriter(Callback):
  201. """
  202. Use VisualDL to log data or image
  203. """
  204. def __init__(self, model):
  205. super(VisualDLWriter, self).__init__(model)
  206. assert six.PY3, "VisualDL requires Python >= 3.5"
  207. try:
  208. from visualdl import LogWriter
  209. except Exception as e:
  210. logger.error('visualdl not found, plaese install visualdl. '
  211. 'for example: `pip install visualdl`.')
  212. raise e
  213. self.vdl_writer = LogWriter(
  214. model.cfg.get('vdl_log_dir', 'vdl_log_dir/scalar'))
  215. self.vdl_loss_step = 0
  216. self.vdl_mAP_step = 0
  217. self.vdl_image_step = 0
  218. self.vdl_image_frame = 0
  219. def on_step_end(self, status):
  220. mode = status['mode']
  221. if dist.get_world_size() < 2 or dist.get_rank() == 0:
  222. if mode == 'train':
  223. training_staus = status['training_staus']
  224. for loss_name, loss_value in training_staus.get().items():
  225. self.vdl_writer.add_scalar(loss_name, loss_value,
  226. self.vdl_loss_step)
  227. self.vdl_loss_step += 1
  228. elif mode == 'test':
  229. ori_image = status['original_image']
  230. result_image = status['result_image']
  231. self.vdl_writer.add_image(
  232. "original/frame_{}".format(self.vdl_image_frame),
  233. ori_image, self.vdl_image_step)
  234. self.vdl_writer.add_image(
  235. "result/frame_{}".format(self.vdl_image_frame),
  236. result_image, self.vdl_image_step)
  237. self.vdl_image_step += 1
  238. # each frame can display ten pictures at most.
  239. if self.vdl_image_step % 10 == 0:
  240. self.vdl_image_step = 0
  241. self.vdl_image_frame += 1
  242. def on_epoch_end(self, status):
  243. mode = status['mode']
  244. if dist.get_world_size() < 2 or dist.get_rank() == 0:
  245. if mode == 'eval':
  246. for metric in self.model._metrics:
  247. for key, map_value in metric.get_results().items():
  248. self.vdl_writer.add_scalar("{}-mAP".format(key),
  249. map_value[0],
  250. self.vdl_mAP_step)
  251. self.vdl_mAP_step += 1
  252. class WandbCallback(Callback):
  253. def __init__(self, model):
  254. super(WandbCallback, self).__init__(model)
  255. try:
  256. import wandb
  257. self.wandb = wandb
  258. except Exception as e:
  259. logger.error('wandb not found, please install wandb. '
  260. 'Use: `pip install wandb`.')
  261. raise e
  262. self.wandb_params = model.cfg.get('wandb', None)
  263. self.save_dir = os.path.join(self.model.cfg.save_dir,
  264. self.model.cfg.filename)
  265. if self.wandb_params is None:
  266. self.wandb_params = {}
  267. for k, v in model.cfg.items():
  268. if k.startswith("wandb_"):
  269. self.wandb_params.update({k.lstrip("wandb_"): v})
  270. self._run = None
  271. if dist.get_world_size() < 2 or dist.get_rank() == 0:
  272. _ = self.run
  273. self.run.config.update(self.model.cfg)
  274. self.run.define_metric("epoch")
  275. self.run.define_metric("eval/*", step_metric="epoch")
  276. self.best_ap = 0
  277. @property
  278. def run(self):
  279. if self._run is None:
  280. if self.wandb.run is not None:
  281. logger.info(
  282. "There is an ongoing wandb run which will be used"
  283. "for logging. Please use `wandb.finish()` to end that"
  284. "if the behaviour is not intended")
  285. self._run = self.wandb.run
  286. else:
  287. self._run = self.wandb.init(**self.wandb_params)
  288. return self._run
  289. def save_model(self,
  290. optimizer,
  291. save_dir,
  292. save_name,
  293. last_epoch,
  294. ema_model=None,
  295. ap=None,
  296. tags=None):
  297. if dist.get_world_size() < 2 or dist.get_rank() == 0:
  298. model_path = os.path.join(save_dir, save_name)
  299. metadata = {}
  300. metadata["last_epoch"] = last_epoch
  301. if ap:
  302. metadata["ap"] = ap
  303. if ema_model is None:
  304. ema_artifact = self.wandb.Artifact(
  305. name="ema_model-{}".format(self.run.id),
  306. type="model",
  307. metadata=metadata)
  308. model_artifact = self.wandb.Artifact(
  309. name="model-{}".format(self.run.id),
  310. type="model",
  311. metadata=metadata)
  312. ema_artifact.add_file(model_path + ".pdema", name="model_ema")
  313. model_artifact.add_file(model_path + ".pdparams", name="model")
  314. self.run.log_artifact(ema_artifact, aliases=tags)
  315. self.run.log_artfact(model_artifact, aliases=tags)
  316. else:
  317. model_artifact = self.wandb.Artifact(
  318. name="model-{}".format(self.run.id),
  319. type="model",
  320. metadata=metadata)
  321. model_artifact.add_file(model_path + ".pdparams", name="model")
  322. self.run.log_artifact(model_artifact, aliases=tags)
  323. def on_step_end(self, status):
  324. mode = status['mode']
  325. if dist.get_world_size() < 2 or dist.get_rank() == 0:
  326. if mode == 'train':
  327. training_status = status['training_staus'].get()
  328. for k, v in training_status.items():
  329. training_status[k] = float(v)
  330. metrics = {"train/" + k: v for k, v in training_status.items()}
  331. self.run.log(metrics)
  332. def on_epoch_end(self, status):
  333. mode = status['mode']
  334. epoch_id = status['epoch_id']
  335. save_name = None
  336. if dist.get_world_size() < 2 or dist.get_rank() == 0:
  337. if mode == 'train':
  338. end_epoch = self.model.cfg.epoch
  339. if (
  340. epoch_id + 1
  341. ) % self.model.cfg.snapshot_epoch == 0 or epoch_id == end_epoch - 1:
  342. save_name = str(
  343. epoch_id
  344. ) if epoch_id != end_epoch - 1 else "model_final"
  345. tags = ["latest", "epoch_{}".format(epoch_id)]
  346. self.save_model(
  347. self.model.optimizer,
  348. self.save_dir,
  349. save_name,
  350. epoch_id + 1,
  351. self.model.use_ema,
  352. tags=tags)
  353. if mode == 'eval':
  354. merged_dict = {}
  355. for metric in self.model._metrics:
  356. for key, map_value in metric.get_results().items():
  357. merged_dict["eval/{}-mAP".format(key)] = map_value[0]
  358. merged_dict["epoch"] = status["epoch_id"]
  359. self.run.log(merged_dict)
  360. if 'save_best_model' in status and status['save_best_model']:
  361. for metric in self.model._metrics:
  362. map_res = metric.get_results()
  363. if 'bbox' in map_res:
  364. key = 'bbox'
  365. elif 'keypoint' in map_res:
  366. key = 'keypoint'
  367. else:
  368. key = 'mask'
  369. if key not in map_res:
  370. logger.warning("Evaluation results empty, this may be due to " \
  371. "training iterations being too few or not " \
  372. "loading the correct weights.")
  373. return
  374. if map_res[key][0] >= self.best_ap:
  375. self.best_ap = map_res[key][0]
  376. save_name = 'best_model'
  377. tags = ["best", "epoch_{}".format(epoch_id)]
  378. self.save_model(
  379. self.model.optimizer,
  380. self.save_dir,
  381. save_name,
  382. last_epoch=epoch_id + 1,
  383. ema_model=self.model.use_ema,
  384. ap=self.best_ap,
  385. tags=tags)
  386. def on_train_end(self, status):
  387. self.run.finish()
  388. class SniperProposalsGenerator(Callback):
  389. def __init__(self, model):
  390. super(SniperProposalsGenerator, self).__init__(model)
  391. ori_dataset = self.model.dataset
  392. self.dataset = self._create_new_dataset(ori_dataset)
  393. self.loader = self.model.loader
  394. self.cfg = self.model.cfg
  395. self.infer_model = self.model.model
  396. def _create_new_dataset(self, ori_dataset):
  397. dataset = copy.deepcopy(ori_dataset)
  398. # init anno_cropper
  399. dataset.init_anno_cropper()
  400. # generate infer roidbs
  401. ori_roidbs = dataset.get_ori_roidbs()
  402. roidbs = dataset.anno_cropper.crop_infer_anno_records(ori_roidbs)
  403. # set new roidbs
  404. dataset.set_roidbs(roidbs)
  405. return dataset
  406. def _eval_with_loader(self, loader):
  407. results = []
  408. with paddle.no_grad():
  409. self.infer_model.eval()
  410. for step_id, data in enumerate(loader):
  411. outs = self.infer_model(data)
  412. for key in ['im_shape', 'scale_factor', 'im_id']:
  413. outs[key] = data[key]
  414. for key, value in outs.items():
  415. if hasattr(value, 'numpy'):
  416. outs[key] = value.numpy()
  417. results.append(outs)
  418. return results
  419. def on_train_end(self, status):
  420. self.loader.dataset = self.dataset
  421. results = self._eval_with_loader(self.loader)
  422. results = self.dataset.anno_cropper.aggregate_chips_detections(results)
  423. # sniper
  424. proposals = []
  425. clsid2catid = {v: k for k, v in self.dataset.catid2clsid.items()}
  426. for outs in results:
  427. batch_res = get_infer_results(outs, clsid2catid)
  428. start = 0
  429. for i, im_id in enumerate(outs['im_id']):
  430. bbox_num = outs['bbox_num']
  431. end = start + bbox_num[i]
  432. bbox_res = batch_res['bbox'][start:end] \
  433. if 'bbox' in batch_res else None
  434. if bbox_res:
  435. proposals += bbox_res
  436. logger.info("save proposals in {}".format(self.cfg.proposals_path))
  437. with open(self.cfg.proposals_path, 'w') as f:
  438. json.dump(proposals, f)