def save_optimizer_weights(model, filepath, overwrite=True, **kwargs):
if not isinstance(filepath, h5py.File):
# If file exists and should not be overwritten.
if not overwrite and os.path.isfile(filepath):
proceed = hdf5_format.ask_to_proceed_with_overwrite(filepath)
if not proceed:
return
f = h5py.File(filepath, mode='w')
opened_new_file = True
else:
f = filepath
opened_new_file = False
try:
model_metadata = saving_utils.model_metadata(
model, include_optimizer=True, require_config=False)
for k, v in model_metadata.items():
if isinstance(v, (dict, list, tuple)):
f.attrs[k] = json.dumps(
v, default=serialization.get_json_type).encode('utf8')
else:
f.attrs[k] = v
if not isinstance(model.optimizer, optimizers.TFOptimizer):
hdf5_format.save_optimizer_weights_to_hdf5_group(f, model.optimizer)
f.flush()
finally:
if opened_new_file:
f.close()
def _save_checkpoint(self, path: pathlib.Path) -> workload.Response:
if not self.is_chief:
return workload.Skipped()
path.mkdir(parents=True, exist_ok=True)
# Save model weights. We use `tf` format because `h5` does not support
# models that subclass `tf.keras.Model` and define custom `call()`
# and/or `train_step()` functions.
self.model.save_weights(str(
path.joinpath("determined-keras-model-weights")),
save_format="tf")
# Save optimizer(s) weights.
with h5py.File(path.joinpath("determined-keras-optimizer-weights.h5"),
"w") as h5file:
for idx, optimizer in enumerate(self.context._optimizers):
opt_group = h5file.create_group(f"optimizer-{idx}")
save_optimizer_weights_to_hdf5_group(opt_group, optimizer)
# Save RNG state.
rng_state = {
"np_rng_state": np.random.get_state(),
"random_rng_state": random.getstate()
}
if version.parse(tf.__version__) < version.parse("2.0.0"):
rng_state["tf_rng_global_seed"] = tf.compat.v1.random.get_seed(
0)[0]
else:
generator = tf.random.get_global_generator()
rng_state["tf2_rng_global_algorithm"] = generator.algorithm
rng_state["tf2_rng_global_state"] = generator.state
with open(path.joinpath("rng_state.pkl"), "wb") as f:
pickle.dump(rng_state, f)
# Save user code.
det.util.write_user_code(path)
# Save callback(s) state.
callbacks_state = self.multiplexer._get_state()
with path.joinpath("determined-callbacks.v1.pkl").open("wb") as f:
pickle.dump(callbacks_state, f)
self.multiplexer._checkpoint_end(path)
return {
"framework": f"tensorflow-{tf.__version__}",
"format": "saved_weights"
}
def _get_optimizer_state(self, optimizer, optimizer_name=None):
state = io.BytesIO()
with h5py.File(state, 'w') as f:
hdf5_format.save_optimizer_weights_to_hdf5_group(f, optimizer)
return state
def save_optimizer_to_hdf5(model, filepath):
with h5py.File(filepath, mode='w') as f:
if (model.optimizer and
not isinstance(model.optimizer, optimizers.TFOptimizer)):
save_optimizer_weights_to_hdf5_group(f, model.optimizer)
f.flush()