def test_build_model(self, block3_strides):
image_size = 321
num_classes = 1000
batch_size = 2
input_shape = (batch_size, image_size, image_size, 3)
model = delf_model.Delf(block3_strides=block3_strides, name='DELF')
model.init_classifiers(num_classes)
images = tf.random.uniform(input_shape,
minval=-1.0,
maxval=1.0,
seed=0)
blocks = {}
# Get global feature by pooling block4 features.
desc_prelogits = model.backbone(images,
intermediates_dict=blocks,
training=False)
desc_logits = model.desc_classification(desc_prelogits)
self.assertAllEqual(desc_prelogits.shape, (batch_size, 2048))
self.assertAllEqual(desc_logits.shape, (batch_size, num_classes))
features = blocks['block3']
attn_prelogits, _, _ = model.attention(features)
attn_logits = model.attn_classification(attn_prelogits)
self.assertAllEqual(attn_prelogits.shape, (batch_size, 1024))
self.assertAllEqual(attn_logits.shape, (batch_size, num_classes))
def __init__(self,
delg_global_features=True,
delg_gem_power=3.0,
delg_embedding_layer_dim=2048,
block3_strides=True,
iou=1.0):
"""Initialization of DELG model.
Args:
delg_global_features: Whether the model uses a DELG-like global feature
head.
delg_gem_power: Power for Generalized Mean pooling in the DELG model. Used
only if 'delg_global_features' is True.
delg_embedding_layer_dim: Size of the FC whitening layer (embedding
layer). Used only if 'delg_global_features' is True.
block3_strides: bool, whether to add strides to the output of block3.
iou: IOU for non-max suppression.
"""
self._stride_factor = 2.0 if block3_strides else 1.0
self._iou = iou
# Setup the DELG model for extraction.
if delg_global_features:
self._model = delg_model.Delg(
block3_strides=block3_strides,
name='DELG',
gem_power=delg_gem_power,
embedding_layer_dim=delg_embedding_layer_dim)
else:
self._model = delf_model.Delf(block3_strides=block3_strides,
name='DELF')
def __init__(self, block3_strides, iou):
"""Initialization of DELF model.
Args:
block3_strides: bool, whether to add strides to the output of block3.
iou: IOU for non-max suppression.
"""
self._stride_factor = 2.0 if block3_strides else 1.0
self._iou = iou
# Setup the DELF model for extraction.
self._model = delf_model.Delf(block3_strides=block3_strides,
name='DELF')
def test_train_step(self, block3_strides):
image_size = 321
num_classes = 1000
batch_size = 2
clip_val = 10.0
input_shape = (batch_size, image_size, image_size, 3)
model = delf_model.Delf(block3_strides=block3_strides, name='DELF')
model.init_classifiers(num_classes)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.001, momentum=0.9)
images = tf.random.uniform(input_shape, minval=0.0, maxval=1.0, seed=0)
labels = tf.random.uniform((batch_size, ),
minval=0,
maxval=model.num_classes - 1,
dtype=tf.int64)
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
def compute_loss(labels, predictions):
per_example_loss = loss_object(labels, predictions)
return tf.nn.compute_average_loss(per_example_loss,
global_batch_size=batch_size)
with tf.GradientTape() as desc_tape:
blocks = {}
desc_prelogits = model.backbone(images,
intermediates_dict=blocks,
training=False)
desc_logits = model.desc_classification(desc_prelogits)
desc_logits = model.desc_classification(desc_prelogits)
desc_loss = compute_loss(labels, desc_logits)
gradients = desc_tape.gradient(desc_loss, model.desc_trainable_weights)
clipped, _ = tf.clip_by_global_norm(gradients, clip_norm=clip_val)
optimizer.apply_gradients(zip(clipped, model.desc_trainable_weights))
with tf.GradientTape() as attn_tape:
block3 = blocks['block3']
block3 = tf.stop_gradient(block3)
attn_prelogits, _, _ = model.attention(block3, training=True)
attn_logits = model.attn_classification(attn_prelogits)
attn_loss = compute_loss(labels, attn_logits)
gradients = attn_tape.gradient(attn_loss, model.attn_trainable_weights)
clipped, _ = tf.clip_by_global_norm(gradients, clip_norm=clip_val)
optimizer.apply_gradients(zip(clipped, model.attn_trainable_weights))
def create_model(num_classes):
"""Define DELF model, and initialize classifiers."""
if FLAGS.delg_global_features:
model = delg_model.Delg(
block3_strides=FLAGS.block3_strides,
name='DELG',
gem_power=FLAGS.delg_gem_power,
embedding_layer_dim=FLAGS.delg_embedding_layer_dim,
scale_factor_init=FLAGS.delg_scale_factor_init,
arcface_margin=FLAGS.delg_arcface_margin)
else:
model = delf_model.Delf(block3_strides=FLAGS.block3_strides,
name='DELF')
model.init_classifiers(num_classes)
return model
def __init__(self, block3_strides, iou):
"""Initialization of DELF model.
Args:
block3_strides: bool, whether to add strides to the output of block3.
iou: IOU for non-max suppression.
"""
self._stride_factor = 2.0 if block3_strides else 1.0
self._iou = iou
# Setup the DELF model for extraction.
self._model = delf_model.Delf(
block3_strides=block3_strides,
name='DELF',
use_dim_reduction=FLAGS.use_autoencoder,
reduced_dimension=FLAGS.autoencoder_dimensions,
dim_expand_channels=FLAGS.local_feature_map_channels)
def test_train_step(self, block3_strides):
image_size = 321
num_classes = 1000
batch_size = 2
clip_val = 10.0
input_shape = (batch_size, image_size, image_size, 3)
model = delf_model.Delf(block3_strides=block3_strides, name='DELF')
model.init_classifiers(num_classes)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.001, momentum=0.9)
images = tf.random.uniform(input_shape, minval=0.0, maxval=1.0, seed=0)
labels = tf.random.uniform((batch_size, ),
minval=0,
maxval=model.num_classes - 1,
dtype=tf.int64)
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
def compute_loss(labels, predictions):
per_example_loss = loss_object(labels, predictions)
return tf.nn.compute_average_loss(per_example_loss,
global_batch_size=batch_size)
with tf.GradientTape() as gradient_tape:
(desc_prelogits, attn_prelogits, _, _, _,
_) = model.global_and_local_forward_pass(images)
# Calculate global loss by applying the descriptor classifier.
desc_logits = model.desc_classification(desc_prelogits)
desc_loss = compute_loss(labels, desc_logits)
# Calculate attention loss by applying the attention block classifier.
attn_logits = model.attn_classification(attn_prelogits)
attn_loss = compute_loss(labels, attn_logits)
# Cumulate global loss and attention loss and backpropagate through the
# descriptor layer and attention layer together.
total_loss = desc_loss + attn_loss
gradients = gradient_tape.gradient(total_loss, model.trainable_weights)
clipped, _ = tf.clip_by_global_norm(gradients, clip_norm=clip_val)
optimizer.apply_gradients(zip(clipped, model.trainable_weights))
def __init__(self,
multi_scale_pool_type='None',
normalize_global_descriptor=False,
input_scales_tensor=None,
delg_global_features=False,
delg_gem_power=3.0,
delg_embedding_layer_dim=2048):
"""Initialization of global feature model.
Args:
multi_scale_pool_type: Type of multi-scale pooling to perform.
normalize_global_descriptor: Whether to L2-normalize global descriptor.
input_scales_tensor: If None, the exported function to be used should be
ExtractFeatures, where an input end-point "input_scales" is added for
the exported model. If not None, the specified 1D tensor of floats will
be hard-coded as the desired input scales, in conjunction with
ExtractFeaturesFixedScales.
delg_global_features: Whether the model uses a DELG-like global feature
head.
delg_gem_power: Power for Generalized Mean pooling in the DELG model. Used
only if 'delg_global_features' is True.
delg_embedding_layer_dim: Size of the FC whitening layer (embedding
layer). Used only if 'delg_global_features' is True.
"""
self._multi_scale_pool_type = multi_scale_pool_type
self._normalize_global_descriptor = normalize_global_descriptor
if input_scales_tensor is None:
self._input_scales_tensor = []
else:
self._input_scales_tensor = input_scales_tensor
# Setup the DELF model for extraction.
if delg_global_features:
self._model = delg_model.Delg(
block3_strides=False,
name='DELG',
gem_power=delg_gem_power,
embedding_layer_dim=delg_embedding_layer_dim)
else:
self._model = delf_model.Delf(block3_strides=False, name='DELF')
def create_model(num_classes):
"""Define DELF model, and initialize classifiers."""
if FLAGS.delg_global_features:
model = delg_model.Delg(
block3_strides=FLAGS.block3_strides,
name='DELG',
gem_power=FLAGS.delg_gem_power,
embedding_layer_dim=FLAGS.delg_embedding_layer_dim,
scale_factor_init=FLAGS.delg_scale_factor_init,
arcface_margin=FLAGS.delg_arcface_margin,
use_dim_reduction=FLAGS.use_autoencoder,
reduced_dimension=FLAGS.autoencoder_dimensions,
dim_expand_channels=FLAGS.local_feature_map_channels)
else:
model = delf_model.Delf(
block3_strides=FLAGS.block3_strides,
name='DELF',
use_dim_reduction=FLAGS.use_autoencoder,
reduced_dimension=FLAGS.autoencoder_dimensions,
dim_expand_channels=FLAGS.local_feature_map_channels)
model.init_classifiers(num_classes)
return model
def __init__(self,
multi_scale_pool_type='None',
normalize_global_descriptor=False,
input_scales_tensor=None):
"""Initialization of global feature model.
Args:
multi_scale_pool_type: Type of multi-scale pooling to perform.
normalize_global_descriptor: Whether to L2-normalize global descriptor.
input_scales_tensor: If None, the exported function to be used should be
ExtractFeatures, where an input end-point "input_scales" is added for
the exported model. If not None, the specified 1D tensor of floats will
be hard-coded as the desired input scales, in conjunction with
ExtractFeaturesFixedScales.
"""
self._multi_scale_pool_type = multi_scale_pool_type
self._normalize_global_descriptor = normalize_global_descriptor
if input_scales_tensor is None:
self._input_scales_tensor = []
else:
self._input_scales_tensor = input_scales_tensor
# Setup the DELF model for extraction.
self._model = delf_model.Delf(block3_strides=False, name='DELF')
def create_model(num_classes):
"""Define DELF model, and initialize classifiers."""
model = delf_model.Delf(block3_strides=FLAGS.block3_strides, name='DELF')
model.init_classifiers(num_classes)
return model
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
export_path = FLAGS.export_path
if os.path.exists(export_path):
raise ValueError('Export_path already exists.')
with tf.Graph().as_default() as g, tf.compat.v1.Session(graph=g) as sess:
# Setup the model for extraction.
model = delf_model.Delf(block3_strides=False, name='DELF')
# Initial forward pass to build model.
images = tf.zeros((1, 321, 321, 3), dtype=tf.float32)
model(images)
# Setup the multiscale extraction.
input_image = tf.compat.v1.placeholder(tf.uint8,
shape=(None, None, 3),
name='input_image')
if FLAGS.input_scales_list is None:
input_scales = tf.compat.v1.placeholder(tf.float32,
shape=[None],
name='input_scales')
else:
input_scales = tf.constant(
[float(s) for s in FLAGS.input_scales_list],
dtype=tf.float32,
shape=[len(FLAGS.input_scales_list)],
name='input_scales')
extracted_features = export_model_utils.ExtractGlobalFeatures(
input_image,
input_scales,
lambda x: model.backbone(x, training=False),
multi_scale_pool_type=FLAGS.multi_scale_pool_type,
normalize_global_descriptor=FLAGS.normalize_global_descriptor)
# Load the weights.
checkpoint_path = FLAGS.ckpt_path
model.load_weights(checkpoint_path)
print('Checkpoint loaded from ', checkpoint_path)
named_input_tensors = {'input_image': input_image}
if FLAGS.input_scales_list is None:
named_input_tensors['input_scales'] = input_scales
# Outputs to the exported model.
named_output_tensors = {}
if FLAGS.multi_scale_pool_type == 'None':
named_output_tensors['global_descriptors'] = tf.identity(
extracted_features, name='global_descriptors')
else:
named_output_tensors['global_descriptor'] = tf.identity(
extracted_features, name='global_descriptor')
# Export the model.
signature_def = (
tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
inputs=_build_tensor_info(named_input_tensors),
outputs=_build_tensor_info(named_output_tensors)))
print('Exporting trained model to:', export_path)
builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(
export_path)
init_op = None
builder.add_meta_graph_and_variables(
sess, [tf.compat.v1.saved_model.tag_constants.SERVING],
signature_def_map={
tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
signature_def
},
main_op=init_op)
builder.save()
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
export_path = FLAGS.export_path
if os.path.exists(export_path):
raise ValueError('Export_path already exists.')
with tf.Graph().as_default() as g, tf.compat.v1.Session(graph=g) as sess:
# Setup the DELF model for extraction.
model = delf_model.Delf(block3_strides=FLAGS.block3_strides,
name='DELF')
# Initial forward pass to build model.
images = tf.zeros((1, 321, 321, 3), dtype=tf.float32)
model(images)
stride_factor = 2.0 if FLAGS.block3_strides else 1.0
# Setup the multiscale keypoint extraction.
input_image = tf.compat.v1.placeholder(tf.uint8,
shape=(None, None, 3),
name='input_image')
input_abs_thres = tf.compat.v1.placeholder(tf.float32,
shape=(),
name='input_abs_thres')
input_scales = tf.compat.v1.placeholder(tf.float32,
shape=[None],
name='input_scales')
input_max_feature_num = tf.compat.v1.placeholder(
tf.int32, shape=(), name='input_max_feature_num')
extracted_features = export_model_utils.ExtractLocalFeatures(
input_image, input_scales, input_max_feature_num, input_abs_thres,
FLAGS.iou, lambda x: model(x, training=False), stride_factor)
# Load the weights.
checkpoint_path = FLAGS.ckpt_path
model.load_weights(checkpoint_path)
print('Checkpoint loaded from ', checkpoint_path)
named_input_tensors = {
'input_image': input_image,
'input_scales': input_scales,
'input_abs_thres': input_abs_thres,
'input_max_feature_num': input_max_feature_num,
}
# Outputs to the exported model.
named_output_tensors = {}
named_output_tensors['boxes'] = tf.identity(extracted_features[0],
name='boxes')
named_output_tensors['features'] = tf.identity(extracted_features[1],
name='features')
named_output_tensors['scales'] = tf.identity(extracted_features[2],
name='scales')
named_output_tensors['scores'] = tf.identity(extracted_features[3],
name='scores')
# Export the model.
signature_def = tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
inputs=_build_tensor_info(named_input_tensors),
outputs=_build_tensor_info(named_output_tensors))
print('Exporting trained model to:', export_path)
builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(
export_path)
init_op = None
builder.add_meta_graph_and_variables(
sess, [tf.compat.v1.saved_model.tag_constants.SERVING],
signature_def_map={
tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
signature_def
},
main_op=init_op)
builder.save()
