dm_learner.py 文件源码

def save_model_for_prediction(self, save_ckpt_fn, vars_to_save=None):
    """Save model data only needed for prediction.

    Args:
      save_ckpt_fn: checkpoint file to save.
      vars_to_save: a list of variables to save.
    """
    if vars_to_save is None:
      vars_to_save = slim.get_model_variables()
      vars_restore_to_exclude = []
      for scope in self.dm_model.restore_scope_exclude:
        vars_restore_to_exclude.extend(slim.get_variables(scope))
      # remove not restored variables.
      vars_to_save = [
          v for v in vars_to_save if v not in vars_restore_to_exclude
      ]
    base_model.save_model(save_ckpt_fn, self.sess, vars_to_save)