-
Notifications
You must be signed in to change notification settings - Fork 422
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add support for saving to and restoring from checkpoints #399
Conversation
…n the same format as tf.train.Checkpoint.
A point for discussion: for my own code, I use anyhow crate to have context in exceptions. Otherwise, one gets an error code without any context. I didn't add anyhow here to avoid add a new dependency to the project, and to be consistent with the rest of this library. But error messages without context are very difficult to diagnose. |
I don't think I'm definitely in favor of including context in errors. We can manually add context in error messages; it's just a matter of doing it. To make it easier, we could add
and use it with
To make this even more concise, we could do the same thing Changes like this to |
Thank you for your reply. Sounds good to me. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks great overall, just a few readability concerns.
src/checkpoint.rs
Outdated
|
||
#[derive(Debug)] | ||
struct SaveRestoreOps { | ||
pub prefix_save: Operation, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There's no need to make these fields pub
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
src/checkpoint.rs
Outdated
/// Add save and restore ops to the graph | ||
fn build_save_ops(&mut self) -> Result<SaveRestoreOps, Status> { | ||
let mut all_variable_ops_opt: Option<Vec<Operation>> = None; | ||
fn make_all_variable_ops(myself: &mut CheckpointMaker) -> Result<Vec<Operation>, Status> { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This should just be a method on CheckpointMaker, since it doesn't capture anything. build_save_ops probably needs to be broken up a bit into smaller methods for readability, anyway.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The function make_all_variable_ops can be removed if we change from variable names to Variable objects. This is what I am going to implement. Then we can go back if desired.
src/option_insert_result.rs
Outdated
@@ -0,0 +1,18 @@ | |||
// Similar to Option<T>.get_or_insert_with, for a function that returns a result. | |||
pub trait OptionInsertWithResult<T> { | |||
fn get_or_insert_with_result<F, E>(&mut self, f: F) -> Result<&T, E> |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Shouldn't this return Result<&mut T, E>
for consistency with get_or_insert_with
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ok, done
src/checkpoint.rs
Outdated
nd.set_attr_type_list("dtypes", &dtypes[..])?; | ||
nd.finish()? | ||
}; | ||
{ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This scope doesn't do anything.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These two scopes (one ending in line 138 and the one starting in 139), that look redundant, were to avoid borrowing errors. The conflict is between the mutable borrow
let mut g = self.scope.graph_mut();
and the immutable borrow
let var_op = self
.scope
.graph()
Because the mutable borrow lasts until the end of the scope.
But it can be done better, by dropping the mutable reference to g with drop(g).
src/checkpoint.rs
Outdated
} | ||
|
||
struct MyScopeData { | ||
pub scope: Scope, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There's no need to make these pub
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done
src/checkpoint.rs
Outdated
pub restore_op: Operation, | ||
} | ||
|
||
/// Checkpointing and restoring struct |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This documentation needs to be a little more in-depth.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ok, I added documentation and a small example at the start of the module.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Since the module is private, the module docs are not visible by default.
Thank you for your review. |
…ects rather than variable names.
Added to commits, the first one Best regards. |
Hello? |
src/checkpoint.rs
Outdated
pub restore_op: Operation, | ||
} | ||
|
||
/// Checkpointing and restoring struct |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Since the module is private, the module docs are not visible by default.
e6f498f
to
9f73a9c
Compare
ping. |
Sorry for the delay, and thank you! |
This pull request contributes support for saving and restoring of variables using tensorflow
SaveV2
/RestoreV2
ops, and therefore it is compatible with checkpoints created withtf.keras.Model.save_weights
.It includes a test case.