improe mdoel saving/laoding

custom_utils/custom_utils.py

@@ -291,22 +291,6 @@ def sparse_to_tensor(L):
     return tf.SparseTensor(idx, data, shape)
 
 
-def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration=''):
-    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, checkpoint
-        assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0', checkpoint
-    if assert_consumed_on:
-        ckpt.restore(manager.latest_checkpoint).assert_consumed()
-    else:
-        if iteration != '':
-            name_base = manager.latest_checkpoint.split('-')[0]
-            ckpt_ind = [i for i in range(len(manager.checkpoints)) if manager.checkpoints[i].endswith(name_base + '-' + iteration)][0]
-            ckpt.restore(manager.checkpoints[ckpt_ind]).assert_existing_objects_matched()#.assert_consumed()
-        else:
-            ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
-
 
 def get_summary_str(model):
     # from torchview import draw_graph