您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

77 行
2.2KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Callback utils
  4. """
  5. class Callbacks:
  6. """"
  7. Handles all registered callbacks for YOLOv5 Hooks
  8. """
  9. # Define the available callbacks
  10. _callbacks = {
  11. 'on_pretrain_routine_start': [],
  12. 'on_pretrain_routine_end': [],
  13. 'on_train_start': [],
  14. 'on_train_epoch_start': [],
  15. 'on_train_batch_start': [],
  16. 'optimizer_step': [],
  17. 'on_before_zero_grad': [],
  18. 'on_train_batch_end': [],
  19. 'on_train_epoch_end': [],
  20. 'on_val_start': [],
  21. 'on_val_batch_start': [],
  22. 'on_val_image_end': [],
  23. 'on_val_batch_end': [],
  24. 'on_val_end': [],
  25. 'on_fit_epoch_end': [], # fit = train + val
  26. 'on_model_save': [],
  27. 'on_train_end': [],
  28. 'teardown': [],
  29. }
  30. def register_action(self, hook, name='', callback=None):
  31. """
  32. Register a new action to a callback hook
  33. Args:
  34. hook The callback hook name to register the action to
  35. name The name of the action for later reference
  36. callback The callback to fire
  37. """
  38. assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
  39. assert callable(callback), f"callback '{callback}' is not callable"
  40. self._callbacks[hook].append({'name': name, 'callback': callback})
  41. def get_registered_actions(self, hook=None):
  42. """"
  43. Returns all the registered actions by callback hook
  44. Args:
  45. hook The name of the hook to check, defaults to all
  46. """
  47. if hook:
  48. return self._callbacks[hook]
  49. else:
  50. return self._callbacks
  51. def run(self, hook, *args, **kwargs):
  52. """
  53. Loop through the registered actions and fire all callbacks
  54. Args:
  55. hook The name of the hook to check, defaults to all
  56. args Arguments to receive from YOLOv5
  57. kwargs Keyword Arguments to receive from YOLOv5
  58. """
  59. assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
  60. for logger in self._callbacks[hook]:
  61. logger['callback'](*args, **kwargs)