Add support for saving to and restoring from checkpoints #399

ramon-garcia · 2023-02-21T02:47:17Z

This pull request contributes support for saving and restoring of variables using tensorflow SaveV2/RestoreV2 ops, and therefore it is compatible with checkpoints created with tf.keras.Model.save_weights.

It includes a test case.

…n the same format as tf.train.Checkpoint.

ramon-garcia · 2023-02-21T12:12:23Z

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.

adamcrume · 2023-03-01T05:13:48Z

I don't think anyhow is appropriate, because it's intended for application code and forces you to use its error type. That would be awkward for a library, and switching to it would be a breaking change (since it would change the return type of functions).

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

impl Status {
  pub fn context(&self, msg: &str) -> Status {
    ...
  }
}

and use it with

my_tensorflow_function().map_err(|s| s.context("foo"))?;

To make this even more concise, we could do the same thing anyhow does and provide an extension trait for Result, but if we were to expose it outside this crate, we'd probably want to call the methods something other than context so they don't conflict for people using both tensorflow and anyhow.

Changes like this to Status should probably be its own PR, though.

ramon-garcia · 2023-03-01T12:12:07Z

Thank you for your reply. Sounds good to me.

adamcrume

Looks great overall, just a few readability concerns.

adamcrume · 2023-03-07T04:32:46Z

src/checkpoint.rs

+
+#[derive(Debug)]
+struct SaveRestoreOps {
+ pub prefix_save: Operation,


There's no need to make these fields pub.

adamcrume · 2023-03-08T03:32:54Z

src/checkpoint.rs

+ /// 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> {


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.

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.

adamcrume · 2023-03-08T03:44:40Z

src/option_insert_result.rs

@@ -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>


Shouldn't this return Result<&mut T, E> for consistency with get_or_insert_with?

adamcrume · 2023-03-08T04:06:00Z

src/checkpoint.rs

+ nd.set_attr_type_list("dtypes", &dtypes[..])?;
+ nd.finish()?
+ };
+ {


This scope doesn't do anything.

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).

adamcrume · 2023-03-08T04:17:09Z

src/checkpoint.rs

+ }
+
+ struct MyScopeData {
+ pub scope: Scope,


There's no need to make these pub.

adamcrume · 2023-03-08T04:37:48Z

src/checkpoint.rs

+ pub restore_op: Operation,
+}
+
+/// Checkpointing and restoring struct


This documentation needs to be a little more in-depth.

Ok, I added documentation and a small example at the start of the module.

Since the module is private, the module docs are not visible by default.

src/checkpoint.rs

ramon-garcia · 2023-03-08T11:56:35Z

Thank you for your review.
In addition, something I am not comfortable with, wouldn't it be better to use an array of variables rather than variable names? When I wrote it, I was concerned about not being able to create a Variable object (for instance, from an imported graph), but it seems to be an independent problem on its own, handled better in its own PR.

…ects rather than variable names.

ramon-garcia · 2023-03-08T15:39:48Z

Added to commits, the first one 013a829. The next one, 2247ed8 changes the specification of variables from strings to Variable objects.

Best regards.

ramon-garcia · 2023-03-20T00:07:42Z

Hello?

src/checkpoint.rs

adamcrume · 2023-03-28T02:58:57Z

src/checkpoint.rs

+ pub restore_op: Operation,
+}
+
+/// Checkpointing and restoring struct


Since the module is private, the module docs are not visible by default.

ramon-garcia · 2023-04-15T05:28:06Z

ping.

adamcrume · 2023-04-19T02:40:53Z

Sorry for the delay, and thank you!

ramon-garcia and others added 2 commits February 21, 2023 03:37

Add support for saving and restoring session state from checkpoints i…

cf74c97

…n the same format as tf.train.Checkpoint.

Merge branch 'tensorflow:master' into checkpoint

43a29fd

ramon-garcia added 4 commits February 22, 2023 14:46

Helper method for option get_or_insert_with_result implemented as trait.

4d2c313

Forgot to make new function public

fee602e

Apply fmt on checkpoint.rs

5c7103a

More fmt fixes

568ff56

adamcrume requested changes Mar 8, 2023

View reviewed changes

ramon-garcia added 2 commits March 8, 2023 15:45

Address reviewer's requests

013a829

In checkpoints, change the specification of variables to Variable obj…

2247ed8

…ects rather than variable names.

ramon-garcia requested a review from adamcrume March 10, 2023 14:58

adamcrume requested changes Mar 28, 2023

View reviewed changes

ramon-garcia force-pushed the checkpoint branch from e6f498f to 9f73a9c Compare March 28, 2023 03:55

Formatting and comment fixes

9f73a9c

Formatting typo in documentation

d04c852

ramon-garcia requested a review from adamcrume April 18, 2023 00:54

adamcrume merged commit 33977d3 into tensorflow:master Apr 19, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for saving to and restoring from checkpoints #399

Add support for saving to and restoring from checkpoints #399

ramon-garcia commented Feb 21, 2023

ramon-garcia commented Feb 21, 2023

adamcrume commented Mar 1, 2023

ramon-garcia commented Mar 1, 2023

adamcrume left a comment

adamcrume Mar 7, 2023

ramon-garcia Mar 8, 2023

adamcrume Mar 8, 2023

ramon-garcia Mar 8, 2023

adamcrume Mar 8, 2023

ramon-garcia Mar 8, 2023

adamcrume Mar 8, 2023

ramon-garcia Mar 8, 2023

adamcrume Mar 8, 2023

ramon-garcia Mar 8, 2023

adamcrume Mar 8, 2023

ramon-garcia Mar 8, 2023

adamcrume Mar 28, 2023

ramon-garcia commented Mar 8, 2023

ramon-garcia commented Mar 8, 2023

ramon-garcia commented Mar 20, 2023

adamcrume Mar 28, 2023

ramon-garcia commented Apr 15, 2023

adamcrume commented Apr 19, 2023

Add support for saving to and restoring from checkpoints #399

Add support for saving to and restoring from checkpoints #399

Conversation

ramon-garcia commented Feb 21, 2023

ramon-garcia commented Feb 21, 2023

adamcrume commented Mar 1, 2023

ramon-garcia commented Mar 1, 2023

adamcrume left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ramon-garcia commented Mar 8, 2023

ramon-garcia commented Mar 8, 2023

ramon-garcia commented Mar 20, 2023

Choose a reason for hiding this comment

ramon-garcia commented Apr 15, 2023

adamcrume commented Apr 19, 2023