callbacks.py 8.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279
  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. import os
  15. import time
  16. import numpy as np
  17. import paddle
  18. from paddle.distributed.parallel import ParallelEnv
  19. from visualdl import LogWriter
  20. from paddlex.paddleseg.utils.progbar import Progbar
  21. import paddlex.paddleseg.utils.logger as logger
  22. class CallbackList(object):
  23. """
  24. Container abstracting a list of callbacks.
  25. Args:
  26. callbacks (list[Callback]): List of `Callback` instances.
  27. """
  28. def __init__(self, callbacks=None):
  29. callbacks = callbacks or []
  30. self.callbacks = [c for c in callbacks]
  31. def append(self, callback):
  32. self.callbacks.append(callback)
  33. def set_params(self, params):
  34. for callback in self.callbacks:
  35. callback.set_params(params)
  36. def set_model(self, model):
  37. for callback in self.callbacks:
  38. callback.set_model(model)
  39. def set_optimizer(self, optimizer):
  40. for callback in self.callbacks:
  41. callback.set_optimizer(optimizer)
  42. def on_iter_begin(self, iter, logs=None):
  43. """Called right before processing a batch.
  44. """
  45. logs = logs or {}
  46. for callback in self.callbacks:
  47. callback.on_iter_begin(iter, logs)
  48. self._t_enter_iter = time.time()
  49. def on_iter_end(self, iter, logs=None):
  50. """Called at the end of a batch.
  51. """
  52. logs = logs or {}
  53. for callback in self.callbacks:
  54. callback.on_iter_end(iter, logs)
  55. self._t_exit_iter = time.time()
  56. def on_train_begin(self, logs=None):
  57. """Called at the beginning of training.
  58. """
  59. logs = logs or {}
  60. for callback in self.callbacks:
  61. callback.on_train_begin(logs)
  62. def on_train_end(self, logs=None):
  63. """Called at the end of training.
  64. """
  65. logs = logs or {}
  66. for callback in self.callbacks:
  67. callback.on_train_end(logs)
  68. def __iter__(self):
  69. return iter(self.callbacks)
  70. class Callback(object):
  71. """Abstract base class used to build new callbacks.
  72. """
  73. def __init__(self):
  74. self.validation_data = None
  75. def set_params(self, params):
  76. self.params = params
  77. def set_model(self, model):
  78. self.model = model
  79. def set_optimizer(self, optimizer):
  80. self.optimizer = optimizer
  81. def on_iter_begin(self, iter, logs=None):
  82. pass
  83. def on_iter_end(self, iter, logs=None):
  84. pass
  85. def on_train_begin(self, logs=None):
  86. pass
  87. def on_train_end(self, logs=None):
  88. pass
  89. class BaseLogger(Callback):
  90. def __init__(self, period=10):
  91. super(BaseLogger, self).__init__()
  92. self.period = period
  93. def _reset(self):
  94. self.totals = {}
  95. def on_train_begin(self, logs=None):
  96. self.totals = {}
  97. def on_iter_end(self, iter, logs=None):
  98. logs = logs or {}
  99. #(iter - 1) // iters_per_epoch + 1
  100. for k, v in logs.items():
  101. if k in self.totals.keys():
  102. self.totals[k] += v
  103. else:
  104. self.totals[k] = v
  105. if iter % self.period == 0 and ParallelEnv().local_rank == 0:
  106. for k in self.totals:
  107. logs[k] = self.totals[k] / self.period
  108. self._reset()
  109. class TrainLogger(Callback):
  110. def __init__(self, log_freq=10):
  111. self.log_freq = log_freq
  112. def _calculate_eta(self, remaining_iters, speed):
  113. if remaining_iters < 0:
  114. remaining_iters = 0
  115. remaining_time = int(remaining_iters * speed)
  116. result = "{:0>2}:{:0>2}:{:0>2}"
  117. arr = []
  118. for i in range(2, -1, -1):
  119. arr.append(int(remaining_time / 60**i))
  120. remaining_time %= 60**i
  121. return result.format(*arr)
  122. def on_iter_end(self, iter, logs=None):
  123. if iter % self.log_freq == 0 and ParallelEnv().local_rank == 0:
  124. total_iters = self.params["total_iters"]
  125. iters_per_epoch = self.params["iters_per_epoch"]
  126. remaining_iters = total_iters - iter
  127. eta = self._calculate_eta(remaining_iters, logs["batch_cost"])
  128. current_epoch = (iter - 1) // self.params["iters_per_epoch"] + 1
  129. loss = logs["loss"]
  130. lr = self.optimizer.get_lr()
  131. batch_cost = logs["batch_cost"]
  132. reader_cost = logs["reader_cost"]
  133. logger.info(
  134. "[TRAIN] epoch={}, iter={}/{}, loss={:.4f}, lr={:.6f}, batch_cost={:.4f}, reader_cost={:.4f} | ETA {}"
  135. .format(current_epoch, iter, total_iters, loss, lr, batch_cost,
  136. reader_cost, eta))
  137. class ProgbarLogger(Callback):
  138. def __init__(self):
  139. super(ProgbarLogger, self).__init__()
  140. def on_train_begin(self, logs=None):
  141. self.verbose = self.params["verbose"]
  142. self.total_iters = self.params["total_iters"]
  143. self.target = self.params["total_iters"]
  144. self.progbar = Progbar(target=self.target, verbose=self.verbose)
  145. self.seen = 0
  146. self.log_values = []
  147. def on_iter_begin(self, iter, logs=None):
  148. #self.seen = 0
  149. if self.seen < self.target:
  150. self.log_values = []
  151. def on_iter_end(self, iter, logs=None):
  152. logs = logs or {}
  153. self.seen += 1
  154. for k in self.params['metrics']:
  155. if k in logs:
  156. self.log_values.append((k, logs[k]))
  157. #if self.verbose and self.seen < self.target and ParallelEnv.local_rank == 0:
  158. #print(self.log_values)
  159. if self.seen < self.target:
  160. self.progbar.update(self.seen, self.log_values)
  161. class ModelCheckpoint(Callback):
  162. def __init__(self,
  163. save_dir,
  164. monitor="miou",
  165. save_best_only=False,
  166. save_params_only=True,
  167. mode="max",
  168. period=1):
  169. super(ModelCheckpoint, self).__init__()
  170. self.monitor = monitor
  171. self.save_dir = save_dir
  172. self.save_best_only = save_best_only
  173. self.save_params_only = save_params_only
  174. self.period = period
  175. self.iters_since_last_save = 0
  176. if mode == "min":
  177. self.monitor_op = np.less
  178. self.best = np.Inf
  179. elif mode == "max":
  180. self.monitor_op = np.greater
  181. self.best = -np.Inf
  182. else:
  183. raise RuntimeError("`mode` is neither \"min\" nor \"max\"!")
  184. def on_train_begin(self, logs=None):
  185. self.verbose = self.params["verbose"]
  186. save_dir = self.save_dir
  187. if not os.path.isdir(save_dir):
  188. if os.path.exists(save_dir):
  189. os.remove(save_dir)
  190. os.makedirs(save_dir)
  191. def on_iter_end(self, iter, logs=None):
  192. logs = logs or {}
  193. self.iters_since_last_save += 1
  194. current_save_dir = os.path.join(self.save_dir, "iter_{}".format(iter))
  195. current_save_dir = os.path.abspath(current_save_dir)
  196. #if self.iters_since_last_save % self.period and ParallelEnv().local_rank == 0:
  197. #self.iters_since_last_save = 0
  198. if iter % self.period == 0 and ParallelEnv().local_rank == 0:
  199. if self.verbose > 0:
  200. print("iter {iter_num}: saving model to {path}".format(
  201. iter_num=iter, path=current_save_dir))
  202. paddle.save(self.model.state_dict(),
  203. os.path.join(current_save_dir, 'model.pdparams'))
  204. if not self.save_params_only:
  205. paddle.save(self.optimizer.state_dict(),
  206. os.path.join(current_save_dir, 'model.pdopt'))
  207. class VisualDL(Callback):
  208. def __init__(self, log_dir="./log", freq=1):
  209. super(VisualDL, self).__init__()
  210. self.log_dir = log_dir
  211. self.freq = freq
  212. def on_train_begin(self, logs=None):
  213. self.writer = LogWriter(self.log_dir)
  214. def on_iter_end(self, iter, logs=None):
  215. logs = logs or {}
  216. if iter % self.freq == 0 and ParallelEnv().local_rank == 0:
  217. for k, v in logs.items():
  218. self.writer.add_scalar("Train/{}".format(k), v, iter)
  219. self.writer.flush()
  220. def on_train_end(self, logs=None):
  221. self.writer.close()