def build_model_base(images, model_name, training, override_params=None):
"""Create a base feature network and return the features before pooling.
Args:
images: input images tensor.
model_name: string, the predefined model name.
training: boolean, whether the model is constructed for training.
override_params: A dictionary of params for overriding. Fields must exist in
efficientnet_model.GlobalParams.
Returns:
features: base features before pooling.
endpoints: the endpoints for each layer.
Raises:
When model_name specified an undefined model, raises NotImplementedError.
When override_params has invalid fields, raises ValueError.
"""
assert isinstance(images, tf.Tensor)
# For backward compatibility.
if override_params and override_params.get('drop_connect_rate', None):
override_params['survival_prob'] = 1 - override_params['drop_connect_rate']
blocks_args, global_params = get_model_params(model_name, override_params)
model = efficientnet_model.Model(blocks_args, global_params, model_name)
outputs = model(images, training=training, features_only=True)
return outputs[0], outputs[1:]
def test_bottleneck_block_with_superpixel_layer(self):
"""Test for creating a model with fused bottleneck block arguments."""
images = tf.zeros((10, 128, 128, 3), dtype=tf.float32)
global_params = efficientnet_model.GlobalParams(
1.0,
1.0,
0,
'channels_last',
num_classes=10,
batch_norm=utils.TpuBatchNormalization)
blocks_args = [
efficientnet_model.BlockArgs(kernel_size=3,
num_repeat=3,
input_filters=3,
output_filters=6,
expand_ratio=6,
id_skip=True,
strides=[2, 2],
conv_type=0,
fused_conv=0,
super_pixel=1)
]
model = efficientnet_model.Model(blocks_args, global_params)
outputs = model(images, training=True)
self.assertEqual((10, 10), outputs[0].shape)
def test_reduction_endpoint_with_single_block_without_sp(self):
"""Test reduction point with single block/layer."""
images = tf.zeros((10, 128, 128, 3), dtype=tf.float32)
global_params = efficientnet_model.GlobalParams(
1.0,
1.0,
0,
'channels_last',
num_classes=10,
batch_norm=utils.TpuBatchNormalization)
blocks_args = [
efficientnet_model.BlockArgs(kernel_size=3,
num_repeat=1,
input_filters=3,
output_filters=6,
expand_ratio=6,
id_skip=False,
strides=[2, 2],
se_ratio=0.8,
conv_type=0,
fused_conv=0,
super_pixel=0)
]
model = efficientnet_model.Model(blocks_args, global_params)
_ = model(images, training=True)
self.assertIn('reduction_1', model.endpoints)
# single block should have one and only one reduction endpoint
self.assertNotIn('reduction_2', model.endpoints)
def test_variables(self):
"""Test for variables in blocks to be included in `model.variables`."""
images = tf.zeros((10, 128, 128, 3), dtype=tf.float32)
global_params = efficientnet_model.GlobalParams(
1.0,
1.0,
0,
'channels_last',
num_classes=10,
batch_norm=utils.TpuBatchNormalization)
blocks_args = [
efficientnet_model.BlockArgs(kernel_size=3,
num_repeat=3,
input_filters=3,
output_filters=6,
expand_ratio=6,
id_skip=False,
strides=[2, 2],
se_ratio=0.8,
conv_type=0,
fused_conv=0,
super_pixel=0)
]
model = efficientnet_model.Model(blocks_args, global_params)
_ = model(images, training=True)
var_names = {var.name for var in model.variables}
self.assertIn('model/blocks_0/conv2d/kernel:0', var_names)
def get_model(model_name, override_params=None, model_dir=None):
"""A helper function to create and return model.
Args:
model_name: string, the predefined model name.
override_params: A dictionary of params for overriding. Fields must exist in
efficientnet_model.GlobalParams.
model_dir: string, optional model dir for saving configs.
Returns:
created model
Raises:
When model_name specified an undefined model, raises NotImplementedError.
When override_params has invalid fields, raises ValueError.
"""
# For backward compatibility.
if override_params and override_params.get('drop_connect_rate', None):
override_params[
'survival_prob'] = 1 - override_params['drop_connect_rate']
if not override_params:
override_params = {}
if model_name.startswith('efficientnet-lite'):
builder = efficientnet_lite_builder
elif model_name.startswith('efficientnet-'):
builder = efficientnet_builder
else:
raise ValueError('Unknown model name {}'.format(model_name))
blocks_args, global_params = builder.get_model_params(
model_name, override_params)
if model_dir:
param_file = os.path.join(model_dir, 'model_params.txt')
if not tf.io.gfile.exists(param_file):
if not tf.io.gfile.exists(model_dir):
tf.io.gfile.mkdir(model_dir)
with tf.io.gfile.GFile(param_file, 'w') as f:
logging.info('writing to %s', param_file)
f.write('model_name= %s\n\n' % model_name)
f.write('global_params= %s\n\n' % str(global_params))
f.write('blocks_args= %s\n\n' % str(blocks_args))
return efficientnet_model.Model(blocks_args, global_params, model_name)
def build_model(images,
model_name,
training,
override_params=None,
model_dir=None,
fine_tuning=False,
features_only=False,
pooled_features_only=False):
"""A helper function to create a model and return predicted logits.
Args:
images: input images tensor.
model_name: string, the predefined model name.
training: boolean, whether the model is constructed for training.
override_params: A dictionary of params for overriding. Fields must exist in
efficientnet_model.GlobalParams.
model_dir: string, optional model dir for saving configs.
fine_tuning: boolean, whether the model is used for finetuning.
features_only: build the base feature network only (excluding final
1x1 conv layer, global pooling, dropout and fc head).
pooled_features_only: build the base network for features extraction (after
1x1 conv layer and global pooling, but before dropout and fc head).
Returns:
logits: the logits tensor of classes.
endpoints: the endpoints for each layer.
Raises:
When model_name specified an undefined model, raises NotImplementedError.
When override_params has invalid fields, raises ValueError.
"""
assert isinstance(images, tf.Tensor)
assert not (features_only and pooled_features_only)
# For backward compatibility.
if override_params and override_params.get('drop_connect_rate', None):
override_params['survival_prob'] = 1 - override_params['drop_connect_rate']
if not training or fine_tuning:
if not override_params:
override_params = {}
override_params['batch_norm'] = utils.BatchNormalization
if fine_tuning:
override_params['relu_fn'] = functools.partial(swish, use_native=False)
blocks_args, global_params = get_model_params(model_name, override_params)
if model_dir:
param_file = os.path.join(model_dir, 'model_params.txt')
if not tf.io.gfile.exists(param_file):
if not tf.io.gfile.exists(model_dir):
tf.io.gfile.makedirs(model_dir)
with tf.io.gfile.GFile(param_file, 'w') as f:
logging.info('writing to %s', param_file)
f.write('model_name= %s\n\n' % model_name)
f.write('global_params= %s\n\n' % str(global_params))
f.write('blocks_args= %s\n\n' % str(blocks_args))
model = efficientnet_model.Model(blocks_args, global_params, model_name)
outputs = model(
images,
training=training,
features_only=features_only,
pooled_features_only=pooled_features_only)
features, endpoints = outputs[0], outputs[1:]
if features_only:
features = tf.identity(features, 'features')
elif pooled_features_only:
features = tf.identity(features, 'pooled_features')
else:
features = tf.identity(features, 'logits')
return features, endpoints