def load_model_from_checkpoint(model_name,
num_classes,
input_size,
ckpt_dir,
quant_aware_train=False):
model = model_builder.create(model_name=model_name,
num_classes=num_classes,
input_size=input_size)
if quant_aware_train:
model = tfmot.quantization.keras.quantize_model(model)
model.summary()
hparams = train_image_classifier.get_default_hparams()
optimizer = train_image_classifier.generate_optimizer(hparams)
loss_fn = train_image_classifier.generate_loss_fn(hparams)
model.compile(optimizer=optimizer, loss=loss_fn, metrics=['accuracy'])
# workaround to fix 'Unresolved object in checkpoint' for optimizer variables
_initialize_model_optimizer(model, input_size, num_classes)
checkpoint_path = os.path.join(ckpt_dir, "ckp")
model.load_weights(checkpoint_path)
return model
def get_model(num_classes, input_size, unfreeze_layers):
model = model_builder.create(
model_name=FLAGS.model_name,
num_classes=num_classes,
input_size=input_size,
unfreeze_layers=unfreeze_layers,
use_coordinates_inputs=FLAGS.use_coordinates_inputs,
base_model_weights=FLAGS.base_model_weights,
seed=FLAGS.random_seed)
return model
def get_model(num_classes):
model = model_builder.create(
model_name=FLAGS.model_name,
num_classes=num_classes,
input_size=FLAGS.input_size,
unfreeze_layers=(FLAGS.unfreeze_layers if FLAGS.fix_resolution else -1),
use_coordinates_inputs=FLAGS.use_coordinates_inputs,
seed=FLAGS.random_seed
)
return model
def _load_model():
model = model_builder.create(
model_name=FLAGS.model_name,
num_classes=FLAGS.num_classes,
input_size=FLAGS.input_size,
use_coordinates_inputs=FLAGS.use_coordinates_inputs,
unfreeze_layers=0)
checkpoint_path = os.path.join(FLAGS.ckpt_dir, "ckp")
model.load_weights(checkpoint_path)
return model
def main(_):
set_random_seeds()
model, base_model = model_builder.create(model_name=FLAGS.model_name,
num_classes=FLAGS.num_classes,
input_size=FLAGS.input_size,
use_coordinates_inputs=FLAGS.use_coordinates_inputs,
unfreeze_layers=0,
return_base_model=True)
checkpoint_path = os.path.join(FLAGS.ckpt_dir, "ckp")
model.load_weights(checkpoint_path)
base_model.save_weights(FLAGS.h5_path, 'h5')
def get_model(num_classes):
model = model_builder.create(model_name=FLAGS.model_name,
num_classes=num_classes,
input_size=FLAGS.input_size,
seed=FLAGS.random_seed)
if FLAGS.ckpt_dir is not None:
checkpoint_path = os.path.join(FLAGS.ckpt_dir, "ckp")
model.load_weights(checkpoint_path)
if FLAGS.quant_aware_train:
model = tfmot.quantization.keras.quantize_model(model)
return model
def _load_model():
model = model_builder.create(
model_name=FLAGS.model_name,
num_classes=FLAGS.num_classes,
input_size=FLAGS.input_size,
use_coordinates_inputs=FLAGS.use_coordinates_inputs,
unfreeze_layers=0)
checkpoint_path = os.path.join(FLAGS.ckpt_dir, "ckp")
model.load_weights(checkpoint_path)
if FLAGS.use_tta:
inputs = [
tf.keras.Input(shape=(FLAGS.input_size, FLAGS.input_size, 3))
for x in range(6)
]
outputs = [model(img_input, training=False) for img_input in inputs]
outputs = tf.keras.layers.Average()(outputs)
tta_model = tf.keras.models.Model(inputs=inputs, outputs=[outputs])
model = tta_model
return model