File size: 2,194 Bytes
8b4c6c7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import inspect
import detectron2.utils.comm as comm
from detectron2.engine import EvalHook as _EvalHook
from detectron2.evaluation.testing import flatten_results_dict
class EvalHook(_EvalHook):
def __init__(self, eval_period, eval_function):
super().__init__(eval_period, eval_function)
func_args = inspect.getfullargspec(eval_function).args
assert {"final_iter", "next_iter"}.issubset(set(func_args)), (
f"Eval function must have either 'final_iter' or 'next_iter' as an argument."
f"Got {func_args} instead."
)
def _do_eval(self, final_iter=False, next_iter=0):
results = self._func(final_iter=final_iter, next_iter=next_iter)
if results:
assert isinstance(
results, dict
), "Eval function must return a dict. Got {} instead.".format(results)
flattened_results = flatten_results_dict(results)
for k, v in flattened_results.items():
try:
v = float(v)
except Exception as e:
raise ValueError(
"[EvalHook] eval_function should return a nested dict of float. "
"Got '{}: {}' instead.".format(k, v)
) from e
self.trainer.storage.put_scalars(**flattened_results, smoothing_hint=False)
# Evaluation may take different time among workers.
# A barrier make them start the next iteration together.
comm.synchronize()
def after_step(self):
next_iter = self.trainer.iter + 1
if self._period > 0 and next_iter % self._period == 0:
# do the last eval in after_train
if next_iter != self.trainer.max_iter:
self._do_eval(next_iter=next_iter)
def after_train(self):
# This condition is to prevent the eval from running after a failed training
if self.trainer.iter + 1 >= self.trainer.max_iter:
self._do_eval(final_iter=True)
# func is likely a closure that holds reference to the trainer
# therefore we clean it to avoid circular reference in the end
del self._func |