return SaverHook()
def save(self, finetune_obj, path, mkdir=True):
if self.values is None:
raise ValueError("No training has been run, cannot save")
folder = os.path.dirname(path)
if not os.path.exists(folder) and mkdir:
os.mkdir(folder)
if self.save_dtype is not None:
self.values = [a.astype(self.save_dtype) for a in self.values]
included, excluded = self.find_trainable_variables()
vars_reduced, vals_reduced = self.remove_unchanged(included, self.values, self.fallback)
var_names = [var.name for var in vars_reduced]
var_dict = dict(zip(var_names, vals_reduced))
assert len(vals_reduced) == len(var_names) == len(var_dict)
joblib.dump((var_dict, finetune_obj), path)
After Change
def save(self, finetune_obj, path, mkdir=True):
ckpt_reader = tf.train.load_checkpoint(finetune_obj.estimator_dir)
variable_map = ckpt_reader.get_variable_to_shape_map()
names = [name for name in variable_map.keys() if self.exclude_matches is None or self.exclude_matches not in name]
names = [name if name.endswith(":0") else name for name in names] // strip the :0 off the end
values = [ckpt_reader.get_tensor(name) for name in names]
names = [name + ":0" for name in names]
