def load_imagenet(ckpt_path):
"""Initialize the network parameters for our xception-lite using ImageNet pretrained weight
Args:
Path to the checkpoint
Returns:
Function that takes a session and initializes the network
# ckpt_path: the full path to the model checkpoint (pre-trained model)
# vars_corresp: A list of `Variable` objects or a dictionary mapping names in the
# checkpoint (pre-trained model) to the corresponding variables to initialize.
"""
reader = tf.train.NewCheckpointReader(ckpt_path)
var_to_shape_map = reader.get_variable_to_shape_map()
vars_corresp = dict()
for v in var_to_shape_map:
if "entry_flow" in v and 'gamma' not in v and 'depthwise/BatchNorm' not in v and 'Momentum' not in v:
vars_corresp[v] = slim.get_model_variables('xfcn/' + v)[0]
elif "middle_flow" in v and 'gamma' not in v and 'depthwise/BatchNorm' not in v and 'Momentum' not in v:
for i in range(1, 3):
if 'unit_{}/'.format(i) in v:
vars_corresp[v] = slim.get_model_variables('xfcn/' + v)[0]
elif "exit_flow" in v and 'gamma' not in v and 'depthwise/BatchNorm' not in v and 'Momentum' not in v:
if 'block1/' in v:
vars_corresp[v] = slim.get_model_variables('xfcn/' + v)[0]
elif 'shortcut' in v and 'Momentum' not in v and 'gamma' not in v:
vars_corresp[v] = slim.get_model_variables('xfcn/' + v)[0]
init_fn = slim.assign_from_checkpoint_fn(ckpt_path, vars_corresp)
return init_fn
def _proc_initializer():
"""Runs initializer for the process. """
logging.info('Proc initializer is called, pid=%i', os.getpid())
config = load_model_proto(FLAGS.detection_pipeline_proto)
model = model_builder.build(config, is_training=False)
image = tf.placeholder(shape=[None, None, 3], dtype=tf.uint8)
preprocessed_inputs, true_image_shapes = model.preprocess(
tf.cast(tf.expand_dims(image, 0), tf.float32))
predictions = model.predict(preprocessed_inputs=preprocessed_inputs,
true_image_shapes=true_image_shapes)
num_proposals = tf.squeeze(predictions['num_proposals'])
proposals = tf.squeeze(predictions['proposal_boxes_normalized'], 0)
proposal_features = tf.reduce_mean(predictions['box_classifier_features'],
[1, 2])
init_fn = slim.assign_from_checkpoint_fn(FLAGS.detection_checkpoint_file,
tf.global_variables())
global tf_session
global tf_inputs
global tf_outputs
config = tf.ConfigProto()
config.allow_soft_placement = True
config.gpu_options.allow_growth = True
tf_session = tf.Session(config=config)
tf_inputs = image
tf_outputs = [num_proposals, proposals, proposal_features]
tf_session.run(tf.global_variables_initializer())
init_fn(tf_session)
uninitialized_variable_names = tf.report_uninitialized_variables()
assert len(tf_session.run(uninitialized_variable_names)) == 0
def _get_init_fn(checkpoint_path, train_dir):
if checkpoint_path is None:
return None
if tf.train.latest_checkpoint(train_dir):
tf.compat.v1.logging.info(
'Ignoring --checkpoint_path because a checkpoint already exists in %s'
% train_dir)
return None
exclusions = []
if checkpoint_exclude_scopes:
exclusions = [
scope.strip() for scope in checkpoint_exclude_scopes.split(',')
]
# TODO(sguada) variables.filter_variables()
variables_to_restore = []
for var in slim.get_model_variables():
for exclusion in exclusions:
if var.op.name.startswith(exclusion):
break
else:
variables_to_restore.append(var)
if tf.io.gfile.isdir(checkpoint_path):
checkpoint_path = tf.train.latest_checkpoint(checkpoint_path)
else:
checkpoint_path = checkpoint_path
tf.compat.v1.logging.info('Fine-tuning from %s' % checkpoint_path)
return slim.assign_from_checkpoint_fn(checkpoint_path,
variables_to_restore,
ignore_missing_vars=True)
def freeze_graph(model,
checkpoint_path,
tensor_shape,
moving_average_decay=0.9999):
"""Converts model ckpts."""
logging.info('Processing ckpt=%s, tensor_shape=%s', checkpoint_path,
tensor_shape)
out_node = 'InceptionV3/Predictions/Reshape_1'
in_node = 'input'
inp = tf.compat.v1.placeholder(shape=[1] + tensor_shape,
dtype=tf.float32,
name=in_node)
_ = model.create(inp, num_classes=3, is_training=False)
ema = tf.train.ExponentialMovingAverage(moving_average_decay)
variables_to_restore = ema.variables_to_restore()
load_ema = slim.assign_from_checkpoint_fn(checkpoint_path,
variables_to_restore,
ignore_missing_vars=True)
with tf.compat.v1.Session() as sess:
sess.run(tf.compat.v1.global_variables_initializer())
load_ema(sess)
graph_def = sess.graph.as_graph_def()
graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
sess, graph_def, [out_node])
graph_def = optimize_for_inference_lib.optimize_for_inference(
graph_def, [in_node], [out_node], tf.float32.as_datatype_enum)
with tf.io.gfile.GFile('model.pb', 'wb') as f:
f.write(graph_def.SerializeToString())
def _get_init_fn():
"""Returns a function to initialize model from a checkpoint."""
if FLAGS.checkpoint_path is None:
return None
# Warn the user if a checkpoint exists in the train_dir. Then we'll be
# ignoring the checkpoint anyway.
if tf.train.latest_checkpoint(FLAGS.train_dir):
tf.logging.info(
'Ignoring --checkpoint_path because a checkpoint already exists in %s',
FLAGS.train_dir)
return None
exclusions = []
if FLAGS.checkpoint_exclude_scopes:
exclusions = [scope.strip()
for scope in FLAGS.checkpoint_exclude_scopes.split(',')]
variables_to_restore = []
for var in slim.get_model_variables():
for exclusion in exclusions:
if var.op.name.startswith(exclusion):
break
else:
variables_to_restore.append(var)
if tf.gfile.IsDirectory(FLAGS.checkpoint_path):
checkpoint_path = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
else:
checkpoint_path = FLAGS.checkpoint_path
tf.logging.info('Fine-tuning from %s', checkpoint_path)
return slim.assign_from_checkpoint_fn(checkpoint_path, variables_to_restore)
def build_deeplabv3(inputs, num_classes, preset_model='DeepLabV3-Res50', weight_decay=1e-5, is_training=True, pretrained_dir="models"):
"""
Builds the DeepLabV3 model.
Arguments:
inputs: The input tensor=
preset_model: Which model you want to use. Select which ResNet model to use for feature extraction
num_classes: Number of classes
Returns:
DeepLabV3 model
"""
if preset_model == 'DeepLabV3-Res50':
with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits, end_points = resnet_v2.resnet_v2_50(inputs, is_training=is_training, scope='resnet_v2_50')
resnet_scope='resnet_v2_50'
# DeepLabV3 requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(os.path.join(pretrained_dir, 'resnet_v2_50.ckpt'), slim.get_model_variables('resnet_v2_50'))
elif preset_model == 'DeepLabV3-Res101':
with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits, end_points = resnet_v2.resnet_v2_101(inputs, is_training=is_training, scope='resnet_v2_101')
resnet_scope='resnet_v2_101'
# DeepLabV3 requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(os.path.join(pretrained_dir, 'resnet_v2_101.ckpt'), slim.get_model_variables('resnet_v2_101'))
elif preset_model == 'DeepLabV3-Res152':
with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits, end_points = resnet_v2.resnet_v2_152(inputs, is_training=is_training, scope='resnet_v2_152')
resnet_scope='resnet_v2_152'
# DeepLabV3 requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(os.path.join(pretrained_dir, 'resnet_v2_152.ckpt'), slim.get_model_variables('resnet_v2_152'))
else:
raise ValueError("Unsupported ResNet model '%s'. This function only supports ResNet 50, ResNet 101, and ResNet 152" % (preset_model))
label_size = tf.shape(inputs)[1:3]
net = AtrousSpatialPyramidPoolingModule(end_points['pool4'])
net = Upsampling(net, label_size)
net = slim.conv2d(net, num_classes, [1, 1], activation_fn=None, scope='logits')
return net, init_fn
def _get_init_fn():
"""Returns a function run by the chief worker to warm-start the training.
Note that the init_fn is only run when initializing the model during the very
first global step.
Returns:
An init function run by the supervisor.
"""
if FLAGS.checkpoint_path is None:
return None
# Warn the user if a checkpoint exists in the train_dir. Then we'll be
# ignoring the checkpoint anyway.
if tf.train.latest_checkpoint(FLAGS.train_dir):
if not FLAGS.continue_training:
raise ValueError(
'continue_training set to False but there is a checkpoint in the training_dir.'
)
tf.compat.v1.logging.info(
'Ignoring --checkpoint_path because a checkpoint already exists in %s'
% FLAGS.train_dir)
return None
exclusions = []
if FLAGS.checkpoint_exclude_scopes:
exclusions = [
scope.strip()
for scope in FLAGS.checkpoint_exclude_scopes.split(',')
]
# TODO(sguada) variables.filter_variables()
variables_to_restore = []
for var in slim.get_model_variables():
excluded = False
for exclusion in exclusions:
if var.op.name.startswith(exclusion):
excluded = True
break
if not excluded:
variables_to_restore.append(var)
if tf.io.gfile.isdir(FLAGS.checkpoint_path):
checkpoint_path = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
else:
checkpoint_path = FLAGS.checkpoint_path
tf.compat.v1.logging.info('Fine-tuning from %s' % checkpoint_path)
return slim.assign_from_checkpoint_fn(
checkpoint_path,
variables_to_restore,
ignore_missing_vars=FLAGS.ignore_missing_vars)
def scaffold_fn():
"""Scaffold function."""
warm_start_hparams = self.hparams.warm_start
if FLAGS.reference_ckpt:
with tf.name_scope('warm_start'):
include_pattern_list = []
if warm_start_hparams.warm_start_encoder:
include_pattern_list.append('ContrastiveModel/Encoder')
if warm_start_hparams.warm_start_projection_head:
include_pattern_list.append(
'ContrastiveModel/ProjectionHead')
if warm_start_hparams.warm_start_classifier:
include_pattern_list.append(
'ContrastiveModel/ClassificationHead')
# This needs to be updated if new optimizers are added.
exclude_pattern_list = [
'Optimizer', 'Momentum', 'RMSProp', 'LARSOptimizer'
]
variables = filter(
lambda v: var_matches_patterns(v,
include_pattern_list),
tf.global_variables())
variables = filter(
lambda v: not var_matches_patterns(
v, exclude_pattern_list), variables)
var_init_fn = slim.assign_from_checkpoint_fn(
tf.train.latest_checkpoint(FLAGS.reference_ckpt),
list(variables),
ignore_missing_vars=(
warm_start_hparams.ignore_missing_checkpoint_vars),
reshape_variables=True)
def init_fn(scaffold, sess):
del scaffold # unused.
if FLAGS.reference_ckpt:
var_init_fn(sess)
return tf.train.Scaffold(init_fn=init_fn)
def get_checkpoint_init_fn():
"""Returns the checkpoint init_fn if the checkpoint is provided."""
if FLAGS.fine_tune_checkpoint:
variables_to_restore = slim.get_variables_to_restore()
global_step_reset = tf.assign(tf.train.get_or_create_global_step(), 0)
# When restoring from a floating point model, the min/max values for
# quantized weights and activations are not present.
# We instruct slim to ignore variables that are missing during restoration
# by setting ignore_missing_vars=True
slim_init_fn = slim.assign_from_checkpoint_fn(
FLAGS.fine_tune_checkpoint,
variables_to_restore,
ignore_missing_vars=True)
def init_fn(sess):
slim_init_fn(sess)
# If we are restoring from a floating point model, we need to initialize
# the global step to zero for the exponential decay to result in
# reasonable learning rates.
sess.run(global_step_reset)
return init_fn
else:
return None
def __init__(self,
network_name,
checkpoint_path,
batch_size,
image_size=None):
self._network_name = network_name
self._checkpoint_path = checkpoint_path
self._batch_size = batch_size
self._image_size = image_size
self._layer = {}
self._global_step = tf.train.get_or_create_global_step()
# Retrieve the function that returns logits and endpoints
self._network_fn = nets_factory.get_network_fn(self._network_name,
num_classes=num_classes,
is_training=False)
# Retrieve the model scope from network factory
self._model_scope = nets_factory.arg_scopes_map[self._network_name]
# Fetch the default image size
self._image_size = self._network_fn.default_image_size
self._filename_queue = tf.FIFOQueue(100000, [tf.string],
shapes=[[]],
name="filename_queue")
self._pl_image_files = tf.placeholder(tf.string,
shape=[None],
name="image_file_list")
self._enqueue_op = self._filename_queue.enqueue_many(
[self._pl_image_files])
self._num_in_queue = self._filename_queue.size()
self._batch_from_queue, self._batch_filenames = self._preproc_image_batch(
self._batch_size, num_threads=4)
#self._image_batch = tf.placeholder_with_default(
# self._batch_from_queue, shape=[self._batch_size, _STRIDE, self._image_size, self._image_size, 3])
self._image_batch = tf.placeholder(
tf.float32, [batch_size, _STRIDE, image_size, image_size, 3])
# Retrieve the logits and network endpoints (for extracting activations)
# Note: endpoints is a dictionary with endpoints[name] = tf.Tensor
self._logits, self._endpoints = self._network_fn(self._image_batch)
# Find the checkpoint file
checkpoint_path = self._checkpoint_path
if tf.gfile.IsDirectory(self._checkpoint_path):
checkpoint_path = tf.train.latest_checkpoint(self._checkpoint_path)
# Load pre-trained weights into the model
variables_to_restore = slim.get_variables_to_restore()
restore_fn = slim.assign_from_checkpoint_fn(self._checkpoint_path,
variables_to_restore)
# Start the session and load the pre-trained weights
self._sess = tf.Session()
restore_fn(self._sess)
# Local variables initializer, needed for queues etc.
self._sess.run(tf.local_variables_initializer())
# Managing the queues and threads
self._coord = tf.train.Coordinator()
self._threads = tf.train.start_queue_runners(coord=self._coord,
sess=self._sess)
def main(unused_argv=None):
with tf.Graph().as_default():
# Force all input processing onto CPU in order to reserve the GPU for the
# forward inference and back-propagation.
device = '/cpu:0' if not FLAGS.ps_tasks else '/job:worker/cpu:0'
with tf.device(
tf.train.replica_device_setter(FLAGS.ps_tasks,
worker_device=device)):
inputs, _ = image_utils.imagenet_inputs(FLAGS.batch_size,
FLAGS.image_size)
# Load style images and select one at random (for each graph execution, a
# new random selection occurs)
_, style_labels, style_gram_matrices = image_utils.style_image_inputs(
os.path.expanduser(FLAGS.style_dataset_file),
batch_size=FLAGS.batch_size,
image_size=FLAGS.image_size,
square_crop=True,
shuffle=True)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Process style and weight flags
num_styles = FLAGS.num_styles
if FLAGS.style_coefficients is None:
style_coefficients = [1.0 for _ in range(num_styles)]
else:
style_coefficients = ast.literal_eval(FLAGS.style_coefficients)
if len(style_coefficients) != num_styles:
raise ValueError(
'number of style coefficients differs from number of styles'
)
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Rescale style weights dynamically based on the current style image
style_coefficient = tf.gather(tf.constant(style_coefficients),
style_labels)
style_weights = dict((key, style_coefficient * value)
for key, value in style_weights.items())
# Define the model
stylized_inputs = model.transform(inputs,
alpha=FLAGS.alpha,
normalizer_params={
'labels': style_labels,
'num_categories': num_styles,
'center': True,
'scale': True
})
# Compute losses.
total_loss, loss_dict = learning.total_loss(
inputs, stylized_inputs, style_gram_matrices, content_weights,
style_weights)
for key, value in loss_dict.items():
tf.summary.scalar(key, value)
instance_norm_vars = [
var for var in slim.get_variables('transformer')
if 'InstanceNorm' in var.name
]
other_vars = [
var for var in slim.get_variables('transformer')
if 'InstanceNorm' not in var.name
]
# Function to restore VGG16 parameters.
init_fn_vgg = slim.assign_from_checkpoint_fn(
vgg.checkpoint_file(), slim.get_variables('vgg_16'))
# Function to restore N-styles parameters.
init_fn_n_styles = slim.assign_from_checkpoint_fn(
os.path.expanduser(FLAGS.checkpoint), other_vars)
def init_fn(session):
init_fn_vgg(session)
init_fn_n_styles(session)
# Set up training.
optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
train_op = slim.learning.create_train_op(
total_loss,
optimizer,
clip_gradient_norm=FLAGS.clip_gradient_norm,
variables_to_train=instance_norm_vars,
summarize_gradients=False)
# Run training.
slim.learning.train(train_op=train_op,
logdir=os.path.expanduser(FLAGS.train_dir),
master=FLAGS.master,
is_chief=FLAGS.task == 0,
number_of_steps=FLAGS.train_steps,
init_fn=init_fn,
save_summaries_secs=FLAGS.save_summaries_secs,
save_interval_secs=FLAGS.save_interval_secs)
model = get_model('inception_v3')
out_node = 'InceptionV3/Predictions/Reshape_1'
in_node = 'input'
inp = tf.compat.v1.placeholder(
shape=[1, args.height, args.width, args.channels],
dtype=tf.float32,
name=in_node)
b = model.create(inp, num_classes=3, is_training=False)
ema = tf.train.ExponentialMovingAverage(args.moving_average_decay)
variables_to_restore = ema.variables_to_restore()
load_ema = slim.assign_from_checkpoint_fn(args.checkpoint,
variables_to_restore,
ignore_missing_vars=True)
with tf.compat.v1.Session() as sess:
sess.run(tf.compat.v1.global_variables_initializer())
load_ema(sess)
graph_def = sess.graph.as_graph_def()
graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
sess, graph_def, [out_node])
graph_def = optimize_for_inference_lib.optimize_for_inference(
graph_def, [in_node], [out_node], tf.float32.as_datatype_enum)
with tf.io.gfile.GFile(args.output, 'wb') as f:
f.write(graph_def.SerializeToString())
def build_frontend(inputs,
frontend,
is_training=True,
pretrained_dir="models",
num_classes=None):
if frontend == 'ResNet50':
with slim.arg_scope(resnet_v2.resnet_arg_scope()):
logits, end_points = resnet_v2.resnet_v2_50(
inputs,
is_training=is_training,
scope='resnet_v2_50',
num_classes=num_classes)
frontend_scope = 'resnet_v2_50'
init_fn = slim.assign_from_checkpoint_fn(
model_path=os.path.join(pretrained_dir, 'resnet_v2_50.ckpt'),
var_list=slim.get_model_variables('resnet_v2_50'),
ignore_missing_vars=True)
elif frontend == 'ResNet101':
with slim.arg_scope(resnet_v2.resnet_arg_scope()):
logits, end_points = resnet_v2.resnet_v2_101(
inputs,
is_training=is_training,
scope='resnet_v2_101',
num_classes=num_classes)
frontend_scope = 'resnet_v2_101'
init_fn = slim.assign_from_checkpoint_fn(
model_path=os.path.join(pretrained_dir, 'resnet_v2_101.ckpt'),
var_list=slim.get_model_variables('resnet_v2_101'),
ignore_missing_vars=True)
elif frontend == 'ResNet152':
with slim.arg_scope(resnet_v2.resnet_arg_scope()):
logits, end_points = resnet_v2.resnet_v2_152(
inputs,
is_training=is_training,
scope='resnet_v2_152',
num_classes=num_classes)
frontend_scope = 'resnet_v2_152'
init_fn = slim.assign_from_checkpoint_fn(
model_path=os.path.join(pretrained_dir, 'resnet_v2_152.ckpt'),
var_list=slim.get_model_variables('resnet_v2_152'),
ignore_missing_vars=True)
elif frontend == 'MobileNetV2':
with slim.arg_scope(mobilenet_v2.training_scope()):
logits, end_points = mobilenet_v2.mobilenet(
inputs,
is_training=is_training,
scope='mobilenet_v2',
base_only=True)
frontend_scope = 'mobilenet_v2'
init_fn = slim.assign_from_checkpoint_fn(
model_path=os.path.join(pretrained_dir, 'mobilenet_v2.ckpt'),
var_list=slim.get_model_variables('mobilenet_v2'),
ignore_missing_vars=True)
elif frontend == 'InceptionV4':
with slim.arg_scope(inception_v4.inception_v4_arg_scope()):
logits, end_points = inception_v4.inception_v4(
inputs, is_training=is_training, scope='inception_v4')
frontend_scope = 'inception_v4'
init_fn = slim.assign_from_checkpoint_fn(
model_path=os.path.join(pretrained_dir, 'inception_v4.ckpt'),
var_list=slim.get_model_variables('inception_v4'),
ignore_missing_vars=True)
else:
raise ValueError(
"Unsupported fronetnd model '%s'. This function only supports ResNet50, ResNet101, ResNet152, and MobileNetV2"
% (frontend))
return logits, end_points, frontend_scope, init_fn
def build_icnet(inputs, label_size, num_classes, preset_model='ICNet', pooling_type = "MAX",
frontend="ResNet101", weight_decay=1e-5, is_training=True, pretrained_dir="models"):
"""
Builds the ICNet model.
Arguments:
inputs: The input tensor
label_size: Size of the final label tensor. We need to know this for proper upscaling
preset_model: Which model you want to use. Select which ResNet model to use for feature extraction
num_classes: Number of classes
pooling_type: Max or Average pooling
Returns:
ICNet model
"""
inputs_4 = tf.image.resize_bilinear(inputs, size=[tf.shape(inputs)[1]*4, tf.shape(inputs)[2]*4])
inputs_2 = tf.image.resize_bilinear(inputs, size=[tf.shape(inputs)[1]*2, tf.shape(inputs)[2]*2])
inputs_1 = inputs
if frontend == 'Res50':
with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits_32, end_points_32 = resnet_v2.resnet_v2_50(inputs_4, is_training=is_training, scope='resnet_v2_50')
logits_16, end_points_16 = resnet_v2.resnet_v2_50(inputs_2, is_training=is_training, scope='resnet_v2_50')
logits_8, end_points_8 = resnet_v2.resnet_v2_50(inputs_1, is_training=is_training, scope='resnet_v2_50')
resnet_scope='resnet_v2_50'
# ICNet requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(os.path.join(pretrained_dir, 'resnet_v2_50.ckpt'), slim.get_model_variables('resnet_v2_50'))
elif frontend == 'Res101':
with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits_32, end_points_32 = resnet_v2.resnet_v2_101(inputs_4, is_training=is_training, scope='resnet_v2_101')
logits_16, end_points_16 = resnet_v2.resnet_v2_101(inputs_2, is_training=is_training, scope='resnet_v2_101')
logits_8, end_points_8 = resnet_v2.resnet_v2_101(inputs_1, is_training=is_training, scope='resnet_v2_101')
resnet_scope='resnet_v2_101'
# ICNet requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(os.path.join(pretrained_dir, 'resnet_v2_101.ckpt'), slim.get_model_variables('resnet_v2_101'))
elif frontend == 'Res152':
with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits_32, end_points_32 = resnet_v2.resnet_v2_152(inputs_4, is_training=is_training, scope='resnet_v2_152')
logits_16, end_points_16 = resnet_v2.resnet_v2_152(inputs_2, is_training=is_training, scope='resnet_v2_152')
logits_8, end_points_8 = resnet_v2.resnet_v2_152(inputs_1, is_training=is_training, scope='resnet_v2_152')
resnet_scope='resnet_v2_152'
# ICNet requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(os.path.join(pretrained_dir, 'resnet_v2_152.ckpt'), slim.get_model_variables('resnet_v2_152'))
else:
raise ValueError("Unsupported ResNet model '%s'. This function only supports ResNet 50, ResNet 101, and ResNet 152" % (frontend))
feature_map_shape = [int(x / 32.0) for x in label_size]
block_32 = PyramidPoolingModule(end_points_32['pool3'], feature_map_shape=feature_map_shape, pooling_type=pooling_type)
out_16, block_16 = CFFBlock(psp_32, end_points_16['pool3'])
out_8, block_8 = CFFBlock(block_16, end_points_8['pool3'])
out_4 = Upsampling_by_scale(out_8, scale=2)
out_4 = slim.conv2d(out_4, num_classes, [1, 1], activation_fn=None)
out_full = Upsampling_by_scale(out_4, scale=2)
out_full = slim.conv2d(out_full, num_classes, [1, 1], activation_fn=None, scope='logits')
net = tf.concat([out_16, out_8, out_4, out_final])
return net, init_fn
def RCNN(inputs, proposals, options, is_training=True):
"""Runs RCNN model on the `inputs`.
Args:
inputs: Input image, a [batch, height, width, 3] uint8 tensor. The pixel
values are in the range of [0, 255].
proposals: Boxes used to crop the image features, using normalized
coordinates. It should be a [batch, max_num_proposals, 4] float tensor
denoting [y1, x1, y2, x2].
options: A fast_rcnn_pb2.FastRCNN proto.
is_training: If true, the model shall be executed in training mode.
Returns:
A [batch, max_num_proposals, feature_dims] tensor.
Raises:
ValueError if options is invalid.
"""
if not isinstance(options, rcnn_pb2.RCNN):
raise ValueError('The options has to be a rcnn_pb2.RCNN proto!')
if inputs.dtype != tf.uint8:
raise ValueError('The inputs has to be a tf.uint8 tensor.')
net_fn = nets_factory.get_network_fn(name=options.feature_extractor_name,
num_classes=1001)
default_image_size = getattr(net_fn, 'default_image_size', 224)
# Preprocess image.
preprocess_fn = preprocessing_factory.get_preprocessing(
options.feature_extractor_name, is_training=False)
inputs = preprocess_fn(inputs,
output_height=None,
output_width=None,
crop_image=False)
# Crop and resize images.
batch = proposals.shape[0]
max_num_proposals = tf.shape(proposals)[1]
box_ind = tf.expand_dims(tf.range(batch), axis=-1)
box_ind = tf.tile(box_ind, [1, max_num_proposals])
cropped_inputs = tf.image.crop_and_resize(
inputs,
boxes=tf.reshape(proposals, [-1, 4]),
box_ind=tf.reshape(box_ind, [-1]),
crop_size=[default_image_size, default_image_size])
# Run CNN.
_, end_points = net_fn(cropped_inputs)
outputs = end_points[options.feature_extractor_endpoint]
outputs = tf.reshape(outputs, [batch, max_num_proposals, -1])
init_fn = slim.assign_from_checkpoint_fn(
options.feature_extractor_checkpoint,
slim.get_model_variables(options.feature_extractor_scope))
def _init_from_ckpt_fn(_, sess):
return init_fn(sess)
return outputs, _init_from_ckpt_fn
def get_init_fn(train_dir=None,
model_checkpoint=None,
exclude_list=None,
include_list=None,
reset_global_step_if_necessary=True,
ignore_missing_vars=True):
"""Gets model initializer function.
The initialization logic is as follows:
1. If a checkpoint is found in `train_dir`, initialize from it.
2. Otherwise, if `model_checkpoint` is valid, initialize from it, and reset
global step if necessary.
3. Otherwise, do not initialize from any checkpoint.
Args:
train_dir: A string as the path to an existing training directory to resume.
Use None to skip.
model_checkpoint: A string as the path to an existing model checkpoint to
initialize from. Use None to skip.
exclude_list: A list of strings for the names of variables not to load.
include_list: A list of strings for the names of variables to load. Use
None to load all variables.
reset_global_step_if_necessary: A boolean for whether to reset global step.
Only used in the case of initializing from an existing checkpoint
`model_checkpoint` rather than resuming training from `train_dir`.
ignore_missing_vars: A boolean for whether to ignore missing variables. If
False, errors will be raised if there is a missing variable.
Returns:
An model initializer function if an existing checkpoint is found. None
otherwise.
"""
# Make sure the exclude list is a list.
if not exclude_list:
exclude_list = []
if train_dir:
train_checkpoint = tf.train.latest_checkpoint(train_dir)
if train_checkpoint:
model_checkpoint = train_checkpoint
logging.info('Resume latest training checkpoint in: %s.',
train_dir)
elif model_checkpoint:
logging.info('Use initial checkpoint: %s.', model_checkpoint)
if reset_global_step_if_necessary:
exclude_list.append('global_step')
logging.info('Reset global step.')
elif model_checkpoint:
logging.info('Use initial checkpoint: %s.', model_checkpoint)
if reset_global_step_if_necessary:
exclude_list.append('global_step')
logging.info('Reset global step.')
if not model_checkpoint:
logging.info('Do not initialize from a checkpoint.')
return None
variables_to_restore = tf_slim.get_variables_to_restore(
include=include_list, exclude=exclude_list)
return tf_slim.assign_from_checkpoint_fn(
model_checkpoint,
variables_to_restore,
ignore_missing_vars=ignore_missing_vars)
def predict(model_root, datasets_dir, model_name, test_image_name):
with tf.Graph().as_default():
tf_global_step = slim.get_or_create_global_step()
test_image = os.path.join(datasets_dir, test_image_name)
# dataset = convert_data.get_datasets('test',dataset_dir=datasets_dir)
network_fn = net_select.get_network_fn(model_name,
num_classes=20,
is_training=False)
batch_size = 1
eval_image_size = network_fn.default_image_size
# images, images_raw, labels = load_batch(datasets_dir,
# height=eval_image_size,
# width=eval_image_size)
image_preprocessing_fn = preprocessing_select.get_preprocessing(
model_name, is_training=False)
image_data = tf.io.read_file(test_image)
image_data = tf.image.decode_jpeg(image_data, channels=3)
image_data = image_preprocessing_fn(image_data, eval_image_size,
eval_image_size)
image_data = tf.expand_dims(image_data, 0)
logits_1, end_points_1 = network_fn(image_data)
attention_maps = tf.reduce_mean(end_points_1['attention_maps'],
axis=-1,
keepdims=True)
attention_maps = tf.image.resize(attention_maps,
[eval_image_size, eval_image_size],
method=tf.image.ResizeMethod.BILINEAR)
bboxes = tf_v1.py_func(mask2bbox, [attention_maps], [tf.float32])
bboxes = tf.reshape(bboxes, [batch_size, 4])
# print(bboxes)
box_ind = tf.range(batch_size, dtype=tf.int32)
images = tf.image.crop_and_resize(
image_data,
bboxes,
box_ind,
crop_size=[eval_image_size, eval_image_size])
logits_2, end_points_2 = network_fn(images, reuse=True)
logits = tf.math.log(
tf.nn.softmax(logits_1) * 0.5 + tf.nn.softmax(logits_2) * 0.5)
checkpoint_path = os.path.join(model_root, model_name)
if tf.io.gfile.isdir(checkpoint_path):
checkpoint_path = tf.train.latest_checkpoint(checkpoint_path)
else:
checkpoint_path = checkpoint_path
init_fn = slim.assign_from_checkpoint_fn(
checkpoint_path, slim.get_variables_to_restore())
# with tf_v1.Session() as sess:
# with slim.queues.QueueRunners(sess):
# sess.run(tf_v1.initialize_local_variables())
# init_fn(sess)
# np_probabilities, np_images_raw, np_labels = sess.run([logits, images_raw, labels])
#
# for i in range(batch_size):
# image = np_images_raw[i, :, :, :]
# true_label = np_labels[i]
# predicted_label = np.argmax(np_probabilities[i, :])
# print('true is {}, predict is {}'.format(true_label, predicted_label))
with tf_v1.Session() as sess:
with slim.queues.QueueRunners(sess):
sess.run(tf_v1.initialize_local_variables())
init_fn(sess)
np_images, np_probabilities = sess.run([image_data, logits])
predicted_label = np.argmax(np_probabilities[0, :])
print(predicted_label)
def main(unused_argv=None):
tf.logging.set_verbosity(tf.logging.INFO)
with tf.Graph().as_default():
# Forces all input processing onto CPU in order to reserve the GPU for the
# forward inference and back-propagation.
device = '/cpu:0' if not FLAGS.ps_tasks else '/job:worker/cpu:0'
with tf.device(
tf.train.replica_device_setter(FLAGS.ps_tasks,
worker_device=device)):
# Loads content images.
content_inputs_, _ = image_utils.imagenet_inputs(
FLAGS.batch_size, FLAGS.image_size)
# Loads style images.
[style_inputs_, _, _] = image_utils.arbitrary_style_image_inputs(
FLAGS.style_dataset_file,
batch_size=FLAGS.batch_size,
image_size=FLAGS.image_size,
shuffle=True,
center_crop=FLAGS.center_crop,
augment_style_images=FLAGS.augment_style_images,
random_style_image_size=FLAGS.random_style_image_size)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Process style and content weight flags.
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Define the model
stylized_images, total_loss, loss_dict, \
_ = build_mobilenet_model.build_mobilenet_model(
content_inputs_,
style_inputs_,
mobilenet_trainable=False,
style_params_trainable=True,
transformer_trainable=True,
mobilenet_end_point='layer_19',
transformer_alpha=FLAGS.alpha,
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=FLAGS.total_variation_weight,
)
# Adding scalar summaries to the tensorboard.
for key in loss_dict:
tf.summary.scalar(key, loss_dict[key])
# Adding Image summaries to the tensorboard.
tf.summary.image('image/0_content_inputs', content_inputs_, 3)
tf.summary.image('image/1_style_inputs_aug', style_inputs_, 3)
tf.summary.image('image/2_stylized_images', stylized_images, 3)
# Set up training
optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
train_op = slim.learning.create_train_op(
total_loss,
optimizer,
clip_gradient_norm=FLAGS.clip_gradient_norm,
summarize_gradients=False)
# Function to restore VGG16 parameters.
init_fn_vgg = slim.assign_from_checkpoint_fn(
vgg.checkpoint_file(), slim.get_variables('vgg_16'))
# Function to restore Mobilenet V2 parameters.
mobilenet_variables_dict = {
var.op.name: var
for var in slim.get_model_variables('MobilenetV2')
}
init_fn_mobilenet = slim.assign_from_checkpoint_fn(
FLAGS.mobilenet_checkpoint, mobilenet_variables_dict)
# Function to restore VGG16 and Mobilenet V2 parameters.
def init_sub_networks(session):
init_fn_vgg(session)
init_fn_mobilenet(session)
# Run training
slim.learning.train(train_op=train_op,
logdir=os.path.expanduser(FLAGS.train_dir),
master=FLAGS.master,
is_chief=FLAGS.task == 0,
number_of_steps=FLAGS.train_steps,
init_fn=init_sub_networks,
save_summaries_secs=FLAGS.save_summaries_secs,
save_interval_secs=FLAGS.save_interval_secs)
def build_gcn(inputs,
num_classes,
preset_model='GCN-Res101',
weight_decay=1e-5,
is_training=True,
upscaling_method="bilinear",
pretrained_dir="models"):
"""
Builds the GCN model.
Arguments:
inputs: The input tensor
preset_model: Which model you want to use. Select which ResNet model to use for feature extraction
num_classes: Number of classes
Returns:
GCN model
"""
if preset_model == 'GCN-Res50':
with slim.arg_scope(
resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits, end_points = resnet_v2.resnet_v2_50(
inputs, is_training=is_training, scope='resnet_v2_50')
resnet_scope = 'resnet_v2_50'
# GCN requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(
os.path.join(pretrained_dir, 'resnet_v2_50.ckpt'),
slim.get_model_variables('resnet_v2_50'))
elif preset_model == 'GCN-Res101':
with slim.arg_scope(
resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits, end_points = resnet_v2.resnet_v2_101(
inputs, is_training=is_training, scope='resnet_v2_101')
resnet_scope = 'resnet_v2_101'
# GCN requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(
os.path.join(pretrained_dir, 'resnet_v2_101.ckpt'),
slim.get_model_variables('resnet_v2_101'))
elif preset_model == 'GCN-Res152':
with slim.arg_scope(
resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits, end_points = resnet_v2.resnet_v2_152(
inputs, is_training=is_training, scope='resnet_v2_152')
resnet_scope = 'resnet_v2_152'
# GCN requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(
os.path.join(pretrained_dir, 'resnet_v2_152.ckpt'),
slim.get_model_variables('resnet_v2_152'))
else:
raise ValueError(
"Unsupported ResNet model '%s'. This function only supports ResNet 101 and ResNet 152"
% (preset_model))
res = [
end_points['pool5'], end_points['pool4'], end_points['pool3'],
end_points['pool2']
]
down_5 = GlobalConvBlock(res[0], n_filters=21, size=3)
down_5 = BoundaryRefinementBlock(down_5, n_filters=21, kernel_size=[3, 3])
down_5 = ConvUpscaleBlock(down_5,
n_filters=21,
kernel_size=[3, 3],
scale=2)
down_4 = GlobalConvBlock(res[1], n_filters=21, size=3)
down_4 = BoundaryRefinementBlock(down_4, n_filters=21, kernel_size=[3, 3])
down_4 = tf.add(down_4, down_5)
down_4 = BoundaryRefinementBlock(down_4, n_filters=21, kernel_size=[3, 3])
down_4 = ConvUpscaleBlock(down_4,
n_filters=21,
kernel_size=[3, 3],
scale=2)
down_3 = GlobalConvBlock(res[2], n_filters=21, size=3)
down_3 = BoundaryRefinementBlock(down_3, n_filters=21, kernel_size=[3, 3])
down_3 = tf.add(down_3, down_4)
down_3 = BoundaryRefinementBlock(down_3, n_filters=21, kernel_size=[3, 3])
down_3 = ConvUpscaleBlock(down_3,
n_filters=21,
kernel_size=[3, 3],
scale=2)
down_2 = GlobalConvBlock(res[3], n_filters=21, size=3)
down_2 = BoundaryRefinementBlock(down_2, n_filters=21, kernel_size=[3, 3])
down_2 = tf.add(down_2, down_3)
down_2 = BoundaryRefinementBlock(down_2, n_filters=21, kernel_size=[3, 3])
down_2 = ConvUpscaleBlock(down_2,
n_filters=21,
kernel_size=[3, 3],
scale=2)
net = BoundaryRefinementBlock(down_2, n_filters=21, kernel_size=[3, 3])
net = ConvUpscaleBlock(net, n_filters=21, kernel_size=[3, 3], scale=2)
net = BoundaryRefinementBlock(net, n_filters=21, kernel_size=[3, 3])
net = slim.conv2d(net,
num_classes, [1, 1],
activation_fn=None,
scope='logits')
return net, init_fn
def __init__(self,
net_name,
snapshot_path,
feature_norm_method=None,
should_restore_classifier=False,
gpu_memory_fraction=None,
vgg_16_heads=None):
"""
Args:
snapshot_path: path or dir with checkpoints
feature_norm_method:
should_restore_classifier: if None - do not restore last layer from the snapshot,
otherwise must be equal to the number of classes of the snapshot.
if vgg_16_heads is not None then the classifiers will be restored anyway.
"""
self.net_name = net_name
if net_name != 'vgg_16_multihead' and vgg_16_heads is not None:
raise ValueError(
'vgg_16_heads must be not None only for vgg_16_multihead')
if net_name == 'vgg_16_multihead' and vgg_16_heads is None:
raise ValueError(
'vgg_16_heads must be not None for vgg_16_multihead')
if tf.io.gfile.isdir(snapshot_path):
snapshot_path = tf.train.latest_checkpoint(snapshot_path)
if not isinstance(feature_norm_method, list):
feature_norm_method = [feature_norm_method]
accepable_methods = [None, 'signed_sqrt', 'unit_norm']
for method in feature_norm_method:
if method not in accepable_methods:
raise ValueError(
'unknown norm method: {}. Use one of {}'.format(
method, accepable_methods))
self.feature_norm_method = feature_norm_method
if vgg_16_heads is not None:
should_restore_classifier = True
if should_restore_classifier:
if vgg_16_heads is None:
reader = pywrap_tensorflow.NewCheckpointReader(snapshot_path)
if net_name == 'inception_v1':
var_value = reader.get_tensor(
'InceptionV1/Logits/Conv2d_0c_1x1/weights')
else:
var_value = reader.get_tensor('vgg_16/fc8/weights')
num_classes = var_value.shape[3]
else:
num_classes = vgg_16_heads
else:
num_classes = 2 if vgg_16_heads is None else vgg_16_heads
network_fn = nets_factory.get_network_fn(net_name,
num_classes=num_classes,
is_training=False)
image_preprocessing_fn = preprocessing_factory.get_preprocessing(
net_name, is_training=False)
eval_image_size = network_fn.default_image_size
self.img_resize_shape = (eval_image_size, eval_image_size
) # (224, 224) for VGG
with tf.Graph().as_default() as graph:
self.graph = graph
with tf.compat.v1.variable_scope('input'):
input_pl = tf.compat.v1.placeholder(
tf.float32,
shape=[None, eval_image_size, eval_image_size, 3],
name='x')
# not used
is_phase_train_pl = tf.compat.v1.placeholder(
tf.bool, shape=tuple(), name='is_phase_train')
function_to_map = lambda x: image_preprocessing_fn(
x, eval_image_size, eval_image_size)
images = tf.map_fn(function_to_map, input_pl)
logits, self.end_points = network_fn(images)
self.__dict__.update(self.end_points)
if net_name == 'inception_v1':
for tensor_name in [
'Branch_0/Conv2d_0a_1x1', 'Branch_1/Conv2d_0a_1x1',
'Branch_1/Conv2d_0b_3x3', 'Branch_2/Conv2d_0a_1x1',
'Branch_2/Conv2d_0b_3x3', 'Branch_3/MaxPool_0a_3x3',
'Branch_3/Conv2d_0b_1x1'
]:
full_tensor_name = 'InceptionV1/InceptionV1/Mixed_4d/' + tensor_name
if 'MaxPool' in tensor_name:
full_tensor_name += '/MaxPool:0'
else:
full_tensor_name += '/Relu:0'
short_name = 'Mixed_4d/' + tensor_name
self.__dict__[short_name] = tf.compat.v1.get_default_graph(
).get_tensor_by_name(full_tensor_name)
self.MaxPool_0a_7x7 = tf.compat.v1.get_default_graph(
).get_tensor_by_name(
"InceptionV1/Logits/MaxPool_0a_7x7/AvgPool:0")
elif net_name in ['vgg_16', 'vgg_16_multihead']:
for layer_name in ['fc6', 'fc7'] + \
['conv{0}/conv{0}_{1}'.format(i, j) for i in range(3, 6) for j in range(1, 4)]:
self.__dict__['vgg_16/{}_prerelu'.format(layer_name)] = \
tf.compat.v1.get_default_graph().get_tensor_by_name("vgg_16/{}/BiasAdd:0".format(layer_name))
config = tf.compat.v1.ConfigProto(
gpu_options=tf.compat.v1.GPUOptions(
per_process_gpu_memory_fraction=gpu_memory_fraction))
self.sess = tf.compat.v1.Session(config=config)
if should_restore_classifier:
variables_to_restore = slim.get_model_variables()
else:
variables_to_restore = [
var for var in slim.get_model_variables()
if not var.op.name.startswith(classifier_scope[net_name])
]
init_fn = slim.assign_from_checkpoint_fn(snapshot_path,
variables_to_restore)
init_fn(self.sess)
gpu_config.gpu_options.per_process_gpu_memory_fraction = 0.8
with tf.Session(config=gpu_config) as sess:
summary_writer = tf.summary.FileWriter('../logs', sess.graph)
sess.run(tf.global_variables_initializer())
train_data.init(sess)
global_step = 0
new_checkpoint = None
if cfg_train_continue:
new_checkpoint = tf.train.latest_checkpoint(
'../checkpoints/checkpoints')
if new_checkpoint:
exclusions = ['global_step']
net_except_logits = slim.get_variables_to_restore(exclude=exclusions)
init_fn = slim.assign_from_checkpoint_fn(new_checkpoint,
net_except_logits,
ignore_missing_vars=True)
init_fn(sess)
print('load params from {}'.format(new_checkpoint))
try:
cur_epoch = 1
while True:
t0 = time.time()
batch_x_img, batch_center_map, batch_scale_map, batch_offset_map, batch_landmark_map = train_data.batch(
sess)
t1 = time.time()
debug_info_, train_loss_, loss_class_, loss_scale_, loss_offset_, loss_landmark_, loss_l2_, summary_, _ = sess.run(
[
debug_info, loss_all_op, loss_class_op, loss_scale_op,
loss_offset_op, loss_landmark_op, loss_l2_op, summary_op,
def build_refinenet(inputs, num_classes, preset_model='RefineNet-Res101', weight_decay=1e-5, is_training=True, upscaling_method="bilinear", pretrained_dir="models"):
"""
Builds the RefineNet model.
Arguments:
inputs: The input tensor
preset_model: Which model you want to use. Select which ResNet model to use for feature extraction
num_classes: Number of classes
Returns:
RefineNet model
"""
if preset_model == 'RefineNet-Res50':
with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits, end_points = resnet_v2.resnet_v2_50(inputs, is_training=is_training, scope='resnet_v2_50')
# RefineNet requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(os.path.join(pretrained_dir, 'resnet_v2_50.ckpt'), slim.get_model_variables('resnet_v2_50'))
elif preset_model == 'RefineNet-Res101':
with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits, end_points = resnet_v2.resnet_v2_101(inputs, is_training=is_training, scope='resnet_v2_101')
# RefineNet requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(os.path.join(pretrained_dir, 'resnet_v2_101.ckpt'), slim.get_model_variables('resnet_v2_101'))
elif preset_model == 'RefineNet-Res152':
with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
logits, end_points = resnet_v2.resnet_v2_152(inputs, is_training=is_training, scope='resnet_v2_152')
# RefineNet requires pre-trained ResNet weights
init_fn = slim.assign_from_checkpoint_fn(os.path.join(pretrained_dir, 'resnet_v2_152.ckpt'), slim.get_model_variables('resnet_v2_152'))
else:
raise ValueError("Unsupported ResNet model '%s'. This function only supports ResNet 101 and ResNet 152" % (preset_model))
high = [end_points['pool5'], end_points['pool4'],
end_points['pool3'], end_points['pool2']]
low = [None, None, None, None]
# Get the feature maps to the proper size with bottleneck
high[0]=slim.conv2d(high[0], 512, 1)
high[1]=slim.conv2d(high[1], 256, 1)
high[2]=slim.conv2d(high[2], 256, 1)
high[3]=slim.conv2d(high[3], 256, 1)
# RefineNet
low[0]=RefineBlock(high_inputs=high[0],low_inputs=None) # Only input ResNet 1/32
low[1]=RefineBlock(high[1],low[0]) # High input = ResNet 1/16, Low input = Previous 1/16
low[2]=RefineBlock(high[2],low[1]) # High input = ResNet 1/8, Low input = Previous 1/8
low[3]=RefineBlock(high[3],low[2]) # High input = ResNet 1/4, Low input = Previous 1/4
# g[3]=Upsampling(g[3],scale=4)
net = low[3]
net = ResidualConvUnit(net)
net = ResidualConvUnit(net)
if upscaling_method.lower() == "conv":
net = ConvUpscaleBlock(net, 128, kernel_size=[3, 3], scale=2)
net = ConvBlock(net, 128)
net = ConvUpscaleBlock(net, 64, kernel_size=[3, 3], scale=2)
net = ConvBlock(net, 64)
elif upscaling_method.lower() == "bilinear":
net = Upsampling(net, scale=4)
net = slim.conv2d(net, num_classes, [1, 1], activation_fn=None, scope='logits')
return net, init_fn
def main(unused_argv=None):
tf.logging.set_verbosity(tf.logging.INFO)
if not tf.gfile.Exists(FLAGS.output_dir):
tf.gfile.MkDir(FLAGS.output_dir)
with tf.Graph().as_default(), tf.Session() as sess:
# Defines place holder for the style image.
style_img_ph = tf.placeholder(tf.float32, shape=[None, None, 3])
if FLAGS.style_square_crop:
style_img_preprocessed = image_utils.center_crop_resize_image(
style_img_ph, FLAGS.style_image_size)
else:
style_img_preprocessed = image_utils.resize_image(style_img_ph,
FLAGS.style_image_size)
# Defines place holder for the content image.
content_img_ph = tf.placeholder(tf.float32, shape=[None, None, 3])
if FLAGS.content_square_crop:
content_img_preprocessed = image_utils.center_crop_resize_image(
content_img_ph, FLAGS.image_size)
else:
content_img_preprocessed = image_utils.resize_image(
content_img_ph, FLAGS.image_size)
# Defines the model.
stylized_images, _, _, bottleneck_feat = build_model.build_model(
content_img_preprocessed,
style_img_preprocessed,
trainable=False,
is_training=False,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
adds_losses=False)
if tf.gfile.IsDirectory(FLAGS.checkpoint):
checkpoint = tf.train.latest_checkpoint(FLAGS.checkpoint)
else:
checkpoint = FLAGS.checkpoint
tf.logging.info('loading latest checkpoint file: {}'.format(checkpoint))
init_fn = slim.assign_from_checkpoint_fn(checkpoint,
slim.get_variables_to_restore())
sess.run([tf.local_variables_initializer()])
init_fn(sess)
# Gets the list of the input style images.
style_img_list = tf.gfile.Glob(FLAGS.style_images_paths)
if len(style_img_list) > FLAGS.maximum_styles_to_evaluate:
np.random.seed(1234)
style_img_list = np.random.permutation(style_img_list)
style_img_list = style_img_list[:FLAGS.maximum_styles_to_evaluate]
# Gets list of input content images.
content_img_list = tf.gfile.Glob(FLAGS.content_images_paths)
for content_i, content_img_path in enumerate(content_img_list):
content_img_np = image_utils.load_np_image_uint8(content_img_path)[:, :, :
3]
content_img_name = os.path.basename(content_img_path)[:-4]
# Saves preprocessed content image.
inp_img_croped_resized_np = sess.run(
content_img_preprocessed, feed_dict={
content_img_ph: content_img_np
})
image_utils.save_np_image(inp_img_croped_resized_np,
os.path.join(FLAGS.output_dir,
'%s.jpg' % (content_img_name)))
# Computes bottleneck features of the style prediction network for the
# identity transform.
identity_params = sess.run(
bottleneck_feat, feed_dict={style_img_ph: content_img_np})
for style_i, style_img_path in enumerate(style_img_list):
if style_i > FLAGS.maximum_styles_to_evaluate:
break
style_img_name = os.path.basename(style_img_path)[:-4]
style_image_np = image_utils.load_np_image_uint8(style_img_path)[:, :, :
3]
if style_i % 10 == 0:
tf.logging.info('Stylizing (%d) %s with (%d) %s' %
(content_i, content_img_name, style_i,
style_img_name))
# Saves preprocessed style image.
style_img_croped_resized_np = sess.run(
style_img_preprocessed, feed_dict={
style_img_ph: style_image_np
})
image_utils.save_np_image(style_img_croped_resized_np,
os.path.join(FLAGS.output_dir,
'%s.jpg' % (style_img_name)))
# Computes bottleneck features of the style prediction network for the
# given style image.
style_params = sess.run(
bottleneck_feat, feed_dict={style_img_ph: style_image_np})
interpolation_weights = ast.literal_eval(FLAGS.interpolation_weights)
# Interpolates between the parameters of the identity transform and
# style parameters of the given style image.
for interp_i, wi in enumerate(interpolation_weights):
stylized_image_res = sess.run(
stylized_images,
feed_dict={
bottleneck_feat:
identity_params * (1 - wi) + style_params * wi,
content_img_ph:
content_img_np
})
# Saves stylized image.
image_utils.save_np_image(
stylized_image_res,
os.path.join(FLAGS.output_dir, '%s_stylized_%s_%d.jpg' %
(content_img_name, style_img_name, interp_i)))
def main(unused_argv=None):
tf.logging.set_verbosity(tf.logging.INFO)
with tf.Graph().as_default():
# Forces all input processing onto CPU in order to reserve the GPU for the
# forward inference and back-propagation.
device = '/cpu:0' if not FLAGS.ps_tasks else '/job:worker/cpu:0'
with tf.device(
tf.train.replica_device_setter(FLAGS.ps_tasks,
worker_device=device)):
# Load content images
content_inputs_, _ = image_utils.imagenet_inputs(
FLAGS.batch_size, FLAGS.image_size)
# Loads style images.
[style_inputs_, _,
style_inputs_orig_] = image_utils.arbitrary_style_image_inputs(
FLAGS.style_dataset_file,
batch_size=FLAGS.batch_size,
image_size=FLAGS.image_size,
shuffle=True,
center_crop=FLAGS.center_crop,
augment_style_images=FLAGS.augment_style_images,
random_style_image_size=FLAGS.random_style_image_size)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Process style and content weight flags.
content_weights = ast.literal_eval(FLAGS.content_weights)
style_weights = ast.literal_eval(FLAGS.style_weights)
# Define the model
stylized_images, \
true_loss, \
_, \
bottleneck_feat = build_mobilenet_model.build_mobilenet_model(
content_inputs_,
style_inputs_,
mobilenet_trainable=True,
style_params_trainable=False,
style_prediction_bottleneck=100,
adds_losses=True,
content_weights=content_weights,
style_weights=style_weights,
total_variation_weight=FLAGS.total_variation_weight,
)
_, inception_bottleneck_feat = build_model.style_prediction(
style_inputs_,
[],
[],
is_training=False,
trainable=False,
inception_end_point='Mixed_6e',
style_prediction_bottleneck=100,
reuse=None,
)
print('PRINTING TRAINABLE VARIABLES')
for x in tf.trainable_variables():
print(x)
mse_loss = tf.losses.mean_squared_error(inception_bottleneck_feat,
bottleneck_feat)
total_loss = mse_loss
if FLAGS.use_true_loss:
true_loss = FLAGS.true_loss_weight * true_loss
total_loss += true_loss
if FLAGS.use_true_loss:
tf.summary.scalar('mse', mse_loss)
tf.summary.scalar('true_loss', true_loss)
tf.summary.scalar('total_loss', total_loss)
tf.summary.image('image/0_content_inputs', content_inputs_, 3)
tf.summary.image('image/1_style_inputs_orig', style_inputs_orig_,
3)
tf.summary.image('image/2_style_inputs_aug', style_inputs_, 3)
tf.summary.image('image/3_stylized_images', stylized_images, 3)
mobilenet_variables_to_restore = slim.get_variables_to_restore(
include=['MobilenetV2'], exclude=['global_step'])
optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
train_op = slim.learning.create_train_op(
total_loss,
optimizer,
clip_gradient_norm=FLAGS.clip_gradient_norm,
summarize_gradients=False)
init_fn = slim.assign_from_checkpoint_fn(
FLAGS.initial_checkpoint,
slim.get_variables_to_restore(
exclude=['MobilenetV2', 'mobilenet_conv', 'global_step']))
init_pretrained_mobilenet = slim.assign_from_checkpoint_fn(
FLAGS.mobilenet_checkpoint, mobilenet_variables_to_restore)
def init_sub_networks(session):
init_fn(session)
init_pretrained_mobilenet(session)
slim.learning.train(train_op=train_op,
logdir=os.path.expanduser(FLAGS.train_dir),
master=FLAGS.master,
is_chief=FLAGS.task == 0,
number_of_steps=FLAGS.train_steps,
init_fn=init_sub_networks,
save_summaries_secs=FLAGS.save_summaries_secs,
save_interval_secs=FLAGS.save_interval_secs)
def main(argv=None):
import os
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
if not tf.gfile.Exists(FLAGS.checkpoint_path):
tf.gfile.MkDir(FLAGS.checkpoint_path)
else:
if not FLAGS.restore:
tf.gfile.DeleteRecursively(FLAGS.checkpoint_path)
tf.gfile.MkDir(FLAGS.checkpoint_path)
input_images = tf.placeholder(tf.float32,
shape=[None, None, None, 3],
name='input_images')
input_score_maps = tf.placeholder(tf.float32,
shape=[None, None, None, 1],
name='input_score_maps')
if FLAGS.geometry == 'RBOX':
input_geo_maps = tf.placeholder(tf.float32,
shape=[None, None, None, 5],
name='input_geo_maps')
else:
input_geo_maps = tf.placeholder(tf.float32,
shape=[None, None, None, 8],
name='input_geo_maps')
input_training_masks = tf.placeholder(tf.float32,
shape=[None, None, None, 1],
name='input_training_masks')
global_step = tf.get_variable('global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
learning_rate = tf.train.exponential_decay(FLAGS.learning_rate,
global_step,
decay_steps=10000,
decay_rate=0.94,
staircase=True)
# add summary
tf.summary.scalar('learning_rate', learning_rate)
opt = tf.train.AdamOptimizer(learning_rate)
# opt = tf.train.MomentumOptimizer(learning_rate, 0.9)
# split
input_images_split = tf.split(input_images, len(gpus))
input_score_maps_split = tf.split(input_score_maps, len(gpus))
input_geo_maps_split = tf.split(input_geo_maps, len(gpus))
input_training_masks_split = tf.split(input_training_masks, len(gpus))
tower_grads = []
reuse_variables = None
for i, gpu_id in enumerate(gpus):
with tf.device('/gpu:%d' % gpu_id):
with tf.name_scope('model_%d' % gpu_id) as scope:
iis = input_images_split[i]
isms = input_score_maps_split[i]
igms = input_geo_maps_split[i]
itms = input_training_masks_split[i]
total_loss, model_loss = tower_loss(iis, isms, igms, itms,
reuse_variables)
batch_norm_updates_op = tf.group(
*tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope))
reuse_variables = True
grads = opt.compute_gradients(total_loss)
tower_grads.append(grads)
grads = average_gradients(tower_grads)
apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)
summary_op = tf.summary.merge_all()
# save moving average
variable_averages = tf.train.ExponentialMovingAverage(
FLAGS.moving_average_decay, global_step)
variables_averages_op = variable_averages.apply(tf.trainable_variables())
# batch norm updates
with tf.control_dependencies(
[variables_averages_op, apply_gradient_op, batch_norm_updates_op]):
train_op = tf.no_op(name='train_op')
saver = tf.train.Saver(tf.global_variables())
summary_writer = tf.summary.FileWriter(FLAGS.checkpoint_path,
tf.get_default_graph())
init = tf.global_variables_initializer()
if FLAGS.pretrained_model_path is not None:
variable_restore_op = slim.assign_from_checkpoint_fn(
FLAGS.pretrained_model_path,
slim.get_trainable_variables(),
ignore_missing_vars=True)
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
if FLAGS.restore:
print('continue training from previous checkpoint')
ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
saver.restore(sess, ckpt)
else:
sess.run(init)
if FLAGS.pretrained_model_path is not None:
variable_restore_op(sess)
data_generator = icdar.get_batch(num_workers=FLAGS.num_readers,
input_size=FLAGS.input_size,
batch_size=FLAGS.batch_size_per_gpu *
len(gpus))
start = time.time()
for step in range(FLAGS.max_steps):
data = next(data_generator)
ml, tl, _ = sess.run(
[model_loss, total_loss, train_op],
feed_dict={
input_images: data[0],
input_score_maps: data[2],
input_geo_maps: data[3],
input_training_masks: data[4]
})
if np.isnan(tl):
print('Loss diverged, stop training')
break
if step % 10 == 0:
avg_time_per_step = (time.time() - start) / 10
avg_examples_per_second = (10 * FLAGS.batch_size_per_gpu *
len(gpus)) / (time.time() - start)
start = time.time()
print(
'Step {:06d}, model loss {:.4f}, total loss {:.4f}, {:.2f} seconds/step, {:.2f} examples/second'
.format(step, ml, tl, avg_time_per_step,
avg_examples_per_second))
if step % FLAGS.save_checkpoint_steps == 0:
saver.save(sess,
FLAGS.checkpoint_path + 'model.ckpt',
global_step=global_step)
if step % FLAGS.save_summary_steps == 0:
_, tl, summary_str = sess.run(
[train_op, total_loss, summary_op],
feed_dict={
input_images: data[0],
input_score_maps: data[2],
input_geo_maps: data[3],
input_training_masks: data[4]
})
summary_writer.add_summary(summary_str, global_step=step)
def main(argv=None):
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
now = datetime.datetime.now()
StyleTime = now.strftime("%Y-%m-%d-%H-%M-%S")
os.makedirs(FLAGS.logs_path + StyleTime)
if not os.path.exists(FLAGS.checkpoint_path):
os.makedirs(FLAGS.checkpoint_path)
input_image = tf.placeholder(tf.float32,
shape=[None, None, None, 3],
name='input_image')
input_bbox = tf.placeholder(tf.float32, shape=[None, 5], name='input_bbox')
input_im_info = tf.placeholder(tf.float32,
shape=[None, 3],
name='input_im_info')
global_step = tf.get_variable('global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
learning_rate = tf.Variable(FLAGS.learning_rate, trainable=False)
tf.summary.scalar('learning_rate', learning_rate)
opt = tf.train.AdamOptimizer(learning_rate)
gpu_id = int(FLAGS.gpu)
with tf.device('/gpu:%d' % gpu_id):
with tf.name_scope('model_%d' % gpu_id) as scope:
bbox_pred, cls_pred, cls_prob = model.model(input_image)
total_loss, model_loss, rpn_cross_entropy, rpn_loss_box = model.loss(
bbox_pred, cls_pred, input_bbox, input_im_info)
batch_norm_updates_op = tf.group(
*tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope))
grads = opt.compute_gradients(total_loss)
apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)
summary_op = tf.summary.merge_all()
variable_averages = tf.train.ExponentialMovingAverage(
FLAGS.moving_average_decay, global_step)
variables_averages_op = variable_averages.apply(tf.trainable_variables())
with tf.control_dependencies(
[variables_averages_op, apply_gradient_op, batch_norm_updates_op]):
train_op = tf.no_op(name='train_op')
saver = tf.train.Saver(tf.global_variables(), max_to_keep=100)
summary_writer = tf.summary.FileWriter(FLAGS.logs_path + StyleTime,
tf.get_default_graph())
init = tf.global_variables_initializer()
if FLAGS.pretrained_model_path is not None:
variable_restore_op = slim.assign_from_checkpoint_fn(
FLAGS.pretrained_model_path,
slim.get_trainable_variables(),
ignore_missing_vars=True)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.95
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
if FLAGS.restore:
ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
restore_step = int(ckpt.split('.')[0].split('_')[-1])
print("continue training from previous checkpoint {}".format(
restore_step))
saver.restore(sess, ckpt)
else:
sess.run(init)
restore_step = 0
if FLAGS.pretrained_model_path is not None:
variable_restore_op(sess)
data_generator = data_provider.get_batch(num_workers=FLAGS.num_readers)
start = time.time()
for step in range(restore_step, FLAGS.max_steps):
data = next(data_generator)
ml, tl, _, summary_str = sess.run(
[model_loss, total_loss, train_op, summary_op],
feed_dict={
input_image: data[0],
input_bbox: data[1],
input_im_info: data[2]
})
summary_writer.add_summary(summary_str, global_step=step)
if step != 0 and step % FLAGS.decay_steps == 0:
sess.run(
tf.assign(learning_rate,
learning_rate.eval() * FLAGS.decay_rate))
if step % 10 == 0:
avg_time_per_step = (time.time() - start) / 10
start = time.time()
print(
'Step {:06d}, model loss {:.4f}, total loss {:.4f}, {:.2f} seconds/step, LR: {:.6f}'
.format(step, ml, tl, avg_time_per_step,
learning_rate.eval()))
if (step + 1) % FLAGS.save_checkpoint_steps == 0:
filename = ('ctpn_{:d}'.format(step + 1) + '.ckpt')
filename = os.path.join(FLAGS.checkpoint_path, filename)
saver.save(sess, filename)
print('Write model to: {:s}'.format(filename))