def load_model(config):
model_checkpoint_path = config.project_config['model_checkpoint_path']
model_name = config.model_name
weights = config.weights
optimizer = Adam(learning_rate=config.train_config['learning_rate'])
optimizer_weights_path = f'{model_path}/optimizer_{checkpoint}.pkl'
with open(optimizer_weights_path, 'rb') as f:
weight_values = pickle.load(f)
optimizer.set_weights(weight_values)
model, loss_inp = DeepICPBuilder(config.net_config).build()
loss = get_loss(config.train_config["loss_alpha"])
source_pts, target_pts, GT = loss_inp
model.add_loss(loss(source_pts, target_pts, GT))
optimizer = Adam(learning_rate=config.train_config['learning_rate'])
model.compile(optimizer=optimizer)
model.load_weights(model_weights)
return model
gen_A_to_B_zeros = [tf.zeros_like(w) for w in gen_A_to_B_vars]
gen_B_to_A_zeros = [tf.zeros_like(w) for w in gen_B_to_A_vars]
disc_A_zeros = [tf.zeros_like(w) for w in disc_A_vars]
disc_B_zeros = [tf.zeros_like(w) for w in disc_B_vars]
# Apply gradients which don't do nothing with Adam
gen_A_to_B_optimizer.apply_gradients(zip(gen_A_to_B_zeros,
gen_A_to_B_vars))
gen_B_to_A_optimizer.apply_gradients(zip(gen_B_to_A_zeros,
gen_B_to_A_vars))
disc_A_optimizer.apply_gradients(zip(disc_A_zeros, disc_A_vars))
disc_B_optimizer.apply_gradients(zip(disc_B_zeros, disc_B_vars))
# Set the weights of the optimizer
gen_A_to_B_optimizer.set_weights(gen_A_to_B_optimizer_weights)
gen_B_to_A_optimizer.set_weights(gen_B_to_A_optimizer_weights)
disc_A_optimizer.set_weights(disc_A_optimizer_weights)
disc_B_optimizer.set_weights(disc_B_optimizer_weights)
# load models again since optimizer might mess them up in first load
gen_A_to_B = tf.keras.models.load_model(
f'models/{DATASET}_{MAX_IMAGE_SIZE}/cycleGAN_e{LOAD_EPOCH:03}_gen_A_to_B'
)
gen_B_to_A = tf.keras.models.load_model(
f'models/{DATASET}_{MAX_IMAGE_SIZE}/cycleGAN_e{LOAD_EPOCH:03}_gen_B_to_A'
)
disc_A = tf.keras.models.load_model(
f'models/{DATASET}_{MAX_IMAGE_SIZE}/cycleGAN_e{LOAD_EPOCH:03}_disc_A')
disc_B = tf.keras.models.load_model(
f'models/{DATASET}_{MAX_IMAGE_SIZE}/cycleGAN_e{LOAD_EPOCH:03}_disc_B')
