File size: 1,976 Bytes
b620cf3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import os.path
import tensorflow as tf
class VisualizeCallback(tf.keras.callbacks.Callback):
def __init__(self, epoch_interval=1, func=lambda model, epoch: None):
super(VisualizeCallback, self).__init__()
self.func = func
self.epoch_interval = epoch_interval
def on_epoch_end(self, epoch, logs=None):
if epoch % self.epoch_interval == 0 and epoch > 0:
self.func(self.model, epoch)
class CheckpointCallback(tf.keras.callbacks.Callback):
def __init__(self, vae, path, epoch_interval=1, restore_training=False, restore_path=None):
super(CheckpointCallback, self).__init__()
self.epoch_interval = epoch_interval
self.path = path
self.vae = vae
self.ckpt = tf.train.Checkpoint(vae=vae,
vae_optimizer=vae.optimizer)
self.ckpt_manager = tf.train.CheckpointManager(checkpoint=self.ckpt,
directory=self.path,
max_to_keep=None)
self.restore_training = restore_training
self.restore_path = restore_path
self._saved = False
def on_epoch_end(self, epoch, logs=None):
if epoch % self.epoch_interval == 0 and epoch > 0:
self.ckpt_manager.save(checkpoint_number=epoch)
def on_train_begin(self, logs=None):
if self.restore_training:
if self.restore_path is None:
self.ckpt.restore(self.ckpt_manager.latest_checkpoint).except_partial()
print("Resume training from checkpoint ", self.ckpt_manager.latest_checkpoint, "\n")
else:
self.ckpt.restore(self.restore_path)
print("resume training from checkpoint ", self.restore_path, "\n")
def on_train_end(self, logs=None):
weights_path = os.path.join(self.path, "trained-vae")
self.ckpt.save(file_prefix=weights_path)
|