Skip to content

best_model_saver

BestModelSaver

Bases: Trace

Save the weights of best model based on a given evaluation metric.

Parameters:

Name Type Description Default
model Union[tf.keras.Model, torch.nn.Module]

A model instance compiled with fe.build.

 required
save_dir str

Folder path into which to save the model.

 required
metric Optional[str]

Eval metric name to monitor. If None, the model's loss will be used.

 None
save_best_mode str

Can be 'min' or 'max'.

 'min'

Raises:

Type Description
AssertionError

If a metric is not provided and it cannot be inferred from the model.
ValueError

If save_best_mode is an unacceptable string.
Source code in fastestimator\fastestimator\trace\io\best_model_saver.py 
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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
class BestModelSaver(Trace):
    """Save the weights of best model based on a given evaluation metric.

    Args:
        model: A model instance compiled with fe.build.
        save_dir: Folder path into which to save the model.
        metric: Eval metric name to monitor. If None, the model's loss will be used.
        save_best_mode: Can be 'min' or 'max'.

    Raises:
        AssertionError: If a `metric` is not provided and it cannot be inferred from the `model`.
        ValueError: If `save_best_mode` is an unacceptable string.
    """
    def __init__(self,
                 model: Union[tf.keras.Model, torch.nn.Module],
                 save_dir: str,
                 metric: Optional[str] = None,
                 save_best_mode: str = "min") -> None:
        if not metric:
            assert hasattr(model, "loss_name"), "cannot infer model loss name, please put the model to UpdateOp first"
            assert len(model.loss_name) == 1, "the model has more than one losses, please provide the metric explicitly"
            metric = next(iter(model.loss_name))
        super().__init__(mode="eval",
                         inputs=metric,
                         outputs=["since_best_{}".format(metric), "{}_{}".format(save_best_mode, metric)])
        self.model = model
        self.save_dir = save_dir
        self.save_best_mode = save_best_mode
        self.since_best = 0
        if self.save_best_mode == "min":
            self.best = np.Inf
            self.monitor_op = np.less
        elif self.save_best_mode == "max":
            self.best = -np.Inf
            self.monitor_op = np.greater
        else:
            raise ValueError("save_best_mode must be either 'min' or 'max'")

    @property
    def metric(self) -> str:
        return self.inputs[0]

    def on_epoch_end(self, data: Data) -> None:
        if self.monitor_op(data[self.metric], self.best):
            self.best = data[self.metric]
            self.since_best = 0
            if self.save_dir:
                model_name = "{}_best_{}".format(self.model.model_name, self.metric)
                model_path = save_model(self.model, self.save_dir, model_name)
                print("FastEstimator-BestModelSaver: Saved model to {}".format(model_path))
        else:
            self.since_best += 1
        data.write_with_log(self.outputs[0], self.since_best)
        data.write_with_log(self.outputs[1], self.best)