Save the model once triggered.
| 13 | |
| 14 | |
| 15 | class ModelSaver(Callback): |
| 16 | """ |
| 17 | Save the model once triggered. |
| 18 | """ |
| 19 | |
| 20 | def __init__(self, max_to_keep=10, |
| 21 | keep_checkpoint_every_n_hours=0.5, |
| 22 | checkpoint_dir=None, |
| 23 | var_collections=None): |
| 24 | """ |
| 25 | Args: |
| 26 | max_to_keep (int): the same as in ``tf.train.Saver``. |
| 27 | keep_checkpoint_every_n_hours (float): the same as in ``tf.train.Saver``. |
| 28 | Note that "keep" does not mean "create", but means "don't delete". |
| 29 | checkpoint_dir (str): Defaults to ``logger.get_logger_dir()``. |
| 30 | var_collections (str or list of str): collection of the variables (or list of collections) to save. |
| 31 | """ |
| 32 | if var_collections is None: |
| 33 | var_collections = [tf.GraphKeys.GLOBAL_VARIABLES] |
| 34 | self._max_to_keep = max_to_keep |
| 35 | self._keep_every_n_hours = keep_checkpoint_every_n_hours |
| 36 | |
| 37 | if not isinstance(var_collections, list): |
| 38 | var_collections = [var_collections] |
| 39 | self.var_collections = var_collections |
| 40 | if checkpoint_dir is None: |
| 41 | checkpoint_dir = logger.get_logger_dir() |
| 42 | if checkpoint_dir is not None: |
| 43 | if not tf.gfile.IsDirectory(checkpoint_dir): # v2: tf.io.gfile.isdir |
| 44 | tf.gfile.MakeDirs(checkpoint_dir) # v2: tf.io.gfile.makedirs |
| 45 | # If None, allow it to be init, but fail later if used |
| 46 | # For example, if chief_only=True, it can still be safely initialized |
| 47 | # in non-chief workers which don't have logger dir |
| 48 | self.checkpoint_dir = fs.normpath(checkpoint_dir) if checkpoint_dir is not None else checkpoint_dir |
| 49 | |
| 50 | def _setup_graph(self): |
| 51 | assert self.checkpoint_dir is not None, \ |
| 52 | "Please provide 'checkpoint_dir' for ModelSaver, or use logger.set_logger_dir()" |
| 53 | vars = [] |
| 54 | for key in self.var_collections: |
| 55 | vars.extend(tf.get_collection(key)) |
| 56 | vars = list(set(vars)) |
| 57 | self.path = os.path.join(self.checkpoint_dir, 'model') |
| 58 | self.saver = tf.train.Saver( |
| 59 | var_list=vars, |
| 60 | max_to_keep=self._max_to_keep, |
| 61 | keep_checkpoint_every_n_hours=self._keep_every_n_hours, |
| 62 | write_version=tf.train.SaverDef.V2, |
| 63 | save_relative_paths=True) |
| 64 | # Scaffold will call saver.build from this collection |
| 65 | tf.add_to_collection(tf.GraphKeys.SAVERS, self.saver) |
| 66 | |
| 67 | def _before_train(self): |
| 68 | # graph is finalized, OK to write it now. |
| 69 | time = datetime.now().strftime('%m%d-%H%M%S') |
| 70 | self.saver.export_meta_graph( |
| 71 | os.path.join(self.checkpoint_dir, |
| 72 | 'graph-{}.meta'.format(time)), |
no outgoing calls
no test coverage detected