luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add testing to static ae

Browse Source
This commit is contained in:
luni
2022-09-25 09:31:50 +00:00
parent 92891f5b26
commit dcda03b942
4 changed files with 29 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
custom_utils/custom_utils.py
+5 -2
View File
@@ -203,10 +203,13 @@ def sparse_to_tensor(L):
return tf.SparseTensor(idx, data, shape)
def load_model(model, checkpoint, disable_assert=False):
def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False):
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
if not disable_assert:
assert manager.latest_checkpoint is not None
assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0'
ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
if assert_consumed_on:
ckpt.restore(manager.latest_checkpoint).assert_consumed()
else:
ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()