def create_pileup_examples(self, dv_call):
"""Creates a tf.Example for DeepVariantCall.
This function calls PileupImageCreator.create_pileup_images on dv_call to
get raw image tensors for each alt_allele option (see docs for details).
These tensors are encoded as pngs, and all of the key information is encoded
as a tf.Example via a call to tf_utils.make_example.
Args:
dv_call: A DeepVariantCall.
Returns:
A list of tf.Example protos.
"""
pileup_images = self.pic.create_pileup_images(dv_call)
if pileup_images is None:
# We cannot build a PileupImage for dv_call, issue a warning.
logging.warning('Could not create PileupImage for candidate at %s:%s',
dv_call.variant.reference_name, dv_call.variant.start)
return []
examples = []
for alt_alleles, image_tensor in pileup_images:
encoded_tensor, shape, tensor_format = self._encode_tensor(image_tensor)
examples.append(
tf_utils.make_example(
dv_call.variant,
alt_alleles,
encoded_tensor,
shape=shape,
image_format=tensor_format))
return examples
def is_valid_call_variants_outputs(call_variants_outputs):
"""Returns True if the call_variants_outputs follows our assumptions.
Args:
call_variants_outputs: list of CallVariantsOutput to check.
Returns:
True if the sanity check passes.
"""
if not call_variants_outputs:
return True # An empty list is a degenerate case.
if not _check_alt_allele_indices(call_variants_outputs):
return False
first_call, other_calls = call_variants_outputs[0], call_variants_outputs[1:]
# Sanity check that all call_variants_outputs have the same `variant`.
for call_to_check in other_calls:
if first_call.variant != call_to_check.variant:
logging.warning('Expected all inputs to merge_predictions to have the '
'same `variant`, but getting %s and %s.',
first_call.variant, call_to_check.variant)
return False
return True
def select_best_match(all_matches):
"""Returns the best LabelerMatch among all_matches.
The best matching LabelerMatch is the one with the lowest match_quality score.
Args:
all_matches: iterable[LabelerMatch]. An iterable of LabelerMatch objects we
want to select the best match from.
Returns:
The best matching LabelerMatch object.
"""
sorted_matches = sorted(all_matches, key=lambda x: x.match_quality)
best = sorted_matches[0]
equivalents = [
f for f in all_matches if f.match_quality == best.match_quality
]
# redacted
if len(equivalents) > 1:
for i, f in enumerate(equivalents):
extra_info = 'best' if i == 0 else i
logging.warning('Equivalent match to best: %s [%s]', f, extra_info)
return equivalents[0]
def _write_checkpoint_metrics(checkpoint_path, metrics_and_values, eval_dir):
"""Writes a JSON of metrics for checkpoint_path in eval_dir.
This function writes out metrics to a JSON for a checkpoint into eval_dir. The
exact path of this file will be computed with:
`checkpoint_metrics_path(checkpoint_path, eval_dir)`
and the values for metrics_and_values (a dict of strings => objects) written
out as key: str(object) into a JSON file.
Args:
checkpoint_path: str; a path to the checkpoint we computed metrics on.
metrics_and_values: dict[string,object]; a dictionary of key/value pairs
containing our metrics. These will be converted to a JSON of key/string
pairs and written out to disk.
eval_dir: str; a path to a directory where we will write out our checkpoint
metrics.
"""
path = checkpoint_metrics_path(checkpoint_path, eval_dir)
serializable = {k: str(v) for k, v in metrics_and_values.iteritems()}
logging.info('Writing checkpoint metrics %s', path)
try:
with tf.gfile.GFile(path, 'w') as fout:
json.dump(serializable, fout, sort_keys=True, indent=4)
except: # pylint: disable=bare-except
# Note we have a bare exception here as as there's no clear TF base
# exception to catch will cover all of the potential issues that might arise
# trying to write our metrics to our metrics file.
logging.warning('Failed to write checkpoint metrics to path %s', path)
def _check_alt_allele_indices(call_variants_outputs):
"""Returns True if and only if the alt allele indices are valid."""
all_alt_allele_indices = sorted([
list(call_variants_output.alt_allele_indices.indices)
for call_variants_output in call_variants_outputs
])
if all_alt_allele_indices != expected_alt_allele_indices(
len(call_variants_outputs[0].variant.alternate_bases)):
logging.warning('Alt allele indices found from call_variants_outputs for '
'variant %s is %s, which is invalid.',
call_variants_outputs[0].variant, all_alt_allele_indices)
return False
return True
def read_training_annotations(training_dir):
"""Reads training data annotations."""
result = []
sub_dirs = tf.gfile.ListDirectory(training_dir)
for sub_dir in sub_dirs:
if not sub_dir.startswith('n'):
logging.warning('Found non-class directory in training dir: %s', sub_dir)
continue
sub_dir_results = read_tiny_imagenet_annotations(
os.path.join(training_dir, sub_dir, sub_dir + '_boxes.txt'),
os.path.join(training_dir, sub_dir, 'images'),
one_label=sub_dir)
result.extend(sub_dir_results)
return result
def _find_matching_variant_in_reader(self, variant):
"""Finds a variant in vcf_reader compatible with variant, if one exists."""
region = variant_utils.variant_position(variant)
matches = [
truth_variant for truth_variant in self._get_truth_variants(region)
if variant.start == truth_variant.start
]
if not matches:
return None
elif len(matches) > 1:
logging.warning(
'Multiple matches detected, keeping first, for variant %s: %s',
variant, matches)
return matches[0]
def _set_target_alignment_info(self, target, alignment):
"""Set target alignment gaps and positions with respect to the reference.
Args:
target: Target, to be processed.
alignment: ssw.AlignmentStructure, target alignment result against the
reference.
Returns:
False, if setting target alignment fails.
Raises:
ValueError: if cigar_op isn't one of the expected ops.
"""
# Note: query_end is inclusive.
if (alignment.query_begin != 0 or
alignment.query_end != len(target.sequence) - 1):
logging.warning('aligner: Target alignment should be end-to-end.')
return False
target.gaps = self._cigar_str_to_gaps(alignment.cigar, 0)
target.ref_pos_mapping = [None] * len(target.sequence)
ref_pos = self.ref_region.start
target_pos = 0
for gap in target.gaps:
pre_gap_len = gap.pos - target_pos
target.ref_pos_mapping[target_pos:gap.pos] = range(
ref_pos, ref_pos + pre_gap_len)
ref_pos += pre_gap_len
target_pos += pre_gap_len
if gap.cigar_op in utils.CIGAR_INSERT_OPS:
target.ref_pos_mapping[target_pos] = ref_pos
target_pos += 1
elif gap.cigar_op in utils.CIGAR_DELETE_OPS:
ref_pos += 1
elif gap.cigar_op in utils.CIGAR_ALIGN_OPS:
target.ref_pos_mapping[target_pos] = ref_pos
ref_pos += 1
target_pos += 1
else:
raise ValueError('Unexpected cigar_op', gap.cigar_op)
pre_gap_len = len(target.sequence) - target_pos
target.ref_pos_mapping[target_pos:] = range(ref_pos, ref_pos + pre_gap_len)
assert target.ref_pos_mapping[0] >= self.ref_region.start
assert target.ref_pos_mapping[-1] <= self.ref_region.end
return True
def _build_metric(metric,
num_categories,
ignored_label,
max_instances_per_category,
intersection_offset=None,
normalize_by_image_size=True):
"""Creates a metric aggregator objet of the given name."""
if metric == 'pq':
logging.warning('One should check Panoptic Quality results against the '
'official COCO API code. Small numerical differences '
'(< 0.1%) can be magnified by rounding.')
return panoptic_quality.PanopticQuality(num_categories, ignored_label,
max_instances_per_category,
intersection_offset)
elif metric == 'pc':
return parsing_covering.ParsingCovering(
num_categories, ignored_label, max_instances_per_category,
intersection_offset, normalize_by_image_size)
else:
raise ValueError('No implementation for metric "%s"' % metric)
def _label_grouped_variants(self, variants):
# redacted
# redacted
# they should be computed in the grouping.
span = ranges.span([variant_utils.variant_range(v) for v in variants])
truths = list(
self._get_truth_variants(
ranges.expand(span, _TRUTH_VARIANTS_QUERY_REGION_EXPANSION_IN_BP)))
if len(truths) > self.max_group_size:
logging.warning(
('Found a large number of variants to label (n_candidates=%d, '
'n_truth=%d) relative to candidate cap of %d. This may make the '
'algorithm very slow.'), len(variants), len(truths),
self.max_group_size)
# redacted
logging.warning('Returning all variants with not-confident markers.')
for variant in variants:
yield variant_labeler.VariantLabel(
is_confident=False, genotype=(-1, -1), variant=variant)
return
ref = self.make_labeler_ref(variants, truths)
labeled_variants = label_variants(variants, truths, ref)
if not labeled_variants:
raise ValueError('Failed to assign labels for variants', variants)
else:
for labeled in labeled_variants:
yield variant_labeler.VariantLabel(
# redacted
# now. Rethink how we establish a variant is confident. Seems like
# it'd be confident if it has a non-ref genotype (as we only
# consider confident truth variants) or if it overlaps the confident
# regions.
is_confident=self._confident_regions.variant_overlaps(labeled),
genotype=tuple(labeled.calls[0].genotype),
variant=labeled)
def _is_thing_array(categories_json, ignored_label):
"""is_thing[category_id] is a bool on if category is "thing" or "stuff"."""
is_thing_dict = {}
for category_json in categories_json:
is_thing_dict[category_json['id']] = bool(category_json['isthing'])
# Check our assumption that the category ids are consecutive.
# Usually metrics should be able to handle this case, but adding a warning
# here.
max_category_id = max(six.iterkeys(is_thing_dict))
if len(is_thing_dict) != max_category_id + 1:
seen_ids = six.viewkeys(is_thing_dict)
all_ids = set(six.moves.range(max_category_id + 1))
unseen_ids = all_ids.difference(seen_ids)
if unseen_ids != {ignored_label}:
logging.warning(
'Nonconsecutive category ids or no category JSON specified for ids: '
'%s', unseen_ids)
is_thing_array = np.zeros(max_category_id + 1)
for category_id, is_thing in six.iteritems(is_thing_dict):
is_thing_array[category_id] = is_thing
return is_thing_array
def make_optimizer_class(cls):
"""Constructs a DP optimizer class from an existing one."""
parent_code = tf.compat.v1.train.Optimizer.compute_gradients.__code__
child_code = cls.compute_gradients.__code__
GATE_OP = tf.compat.v1.train.Optimizer.GATE_OP #pylint: disable=invalid-name
if child_code is not parent_code:
logging.warning(
'WARNING: Calling make_optimizer_class() on class %s that overrides '
'method compute_gradients(). Check to ensure that '
'make_optimizer_class() does not interfere with overridden version.',
cls.__name__)
class DPOptimizerClass(cls):
"""Differentially private subclass of given class cls."""
def __init__(
self,
dp_sum_query,
num_microbatches=None,
unroll_microbatches=False,
*args, # pylint: disable=keyword-arg-before-vararg, g-doc-args
**kwargs):
"""Initialize the DPOptimizerClass.
Args:
dp_sum_query: DPQuery object, specifying differential privacy
mechanism to use.
num_microbatches: How many microbatches into which the minibatch is
split. If None, will default to the size of the minibatch, and
per-example gradients will be computed.
unroll_microbatches: If true, processes microbatches within a Python
loop instead of a tf.while_loop. Can be used if using a tf.while_loop
raises an exception.
"""
super(DPOptimizerClass, self).__init__(*args, **kwargs)
self._dp_sum_query = dp_sum_query
self._num_microbatches = num_microbatches
self._global_state = self._dp_sum_query.initial_global_state()
# TODO(b/122613513): Set unroll_microbatches=True to avoid this bug.
# Beware: When num_microbatches is large (>100), enabling this parameter
# may cause an OOM error.
self._unroll_microbatches = unroll_microbatches
def compute_gradients(self,
loss,
var_list,
gate_gradients=GATE_OP,
aggregation_method=None,
colocate_gradients_with_ops=False,
grad_loss=None,
gradient_tape=None):
if callable(loss):
# TF is running in Eager mode, check we received a vanilla tape.
if not gradient_tape:
raise ValueError(
'When in Eager mode, a tape needs to be passed.')
vector_loss = loss()
if self._num_microbatches is None:
self._num_microbatches = tf.shape(input=vector_loss)[0]
sample_state = self._dp_sum_query.initial_sample_state(
var_list)
microbatches_losses = tf.reshape(vector_loss,
[self._num_microbatches, -1])
sample_params = (self._dp_sum_query.derive_sample_params(
self._global_state))
def process_microbatch(i, sample_state):
"""Process one microbatch (record) with privacy helper."""
microbatch_loss = tf.reduce_mean(
input_tensor=tf.gather(microbatches_losses, [i]))
grads = gradient_tape.gradient(microbatch_loss, var_list)
sample_state = self._dp_sum_query.accumulate_record(
sample_params, sample_state, grads)
return sample_state
for idx in range(self._num_microbatches):
sample_state = process_microbatch(idx, sample_state)
grad_sums, self._global_state = (
self._dp_sum_query.get_noised_result(
sample_state, self._global_state))
def normalize(v):
return v / tf.cast(self._num_microbatches, tf.float32)
final_grads = tf.nest.map_structure(normalize, grad_sums)
grads_and_vars = list(zip(final_grads, var_list))
return grads_and_vars
else:
# TF is running in graph mode, check we did not receive a gradient tape.
if gradient_tape:
raise ValueError(
'When in graph mode, a tape should not be passed.')
# Note: it would be closer to the correct i.i.d. sampling of records if
# we sampled each microbatch from the appropriate binomial distribution,
# although that still wouldn't be quite correct because it would be
# sampling from the dataset without replacement.
if self._num_microbatches is None:
self._num_microbatches = tf.shape(input=loss)[0]
microbatches_losses = tf.reshape(loss,
[self._num_microbatches, -1])
sample_params = (self._dp_sum_query.derive_sample_params(
self._global_state))
def process_microbatch(i, sample_state):
"""Process one microbatch (record) with privacy helper."""
grads, _ = zip(*super(cls, self).compute_gradients(
tf.reduce_mean(
input_tensor=tf.gather(microbatches_losses, [i])),
var_list, gate_gradients, aggregation_method,
colocate_gradients_with_ops, grad_loss))
grads_list = [
g if g is not None else tf.zeros_like(v)
for (g, v) in zip(list(grads), var_list)
]
sample_state = self._dp_sum_query.accumulate_record(
sample_params, sample_state, grads_list)
return sample_state
if var_list is None:
var_list = (tf.compat.v1.trainable_variables(
) + tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.TRAINABLE_RESOURCE_VARIABLES))
sample_state = self._dp_sum_query.initial_sample_state(
var_list)
print(var_list)
if self._unroll_microbatches:
for idx in range(self._num_microbatches):
sample_state = process_microbatch(idx, sample_state)
else:
# Use of while_loop here requires that sample_state be a nested
# structure of tensors. In general, we would prefer to allow it to be
# an arbitrary opaque type.
cond_fn = lambda i, _: tf.less(i, self._num_microbatches)
body_fn = lambda i, state: [tf.add(i, 1), process_microbatch(i, state)] # pylint: disable=line-too-long
idx = tf.constant(0)
_, sample_state = tf.while_loop(
cond=cond_fn,
body=body_fn,
loop_vars=[idx, sample_state])
grad_sums, self._global_state = (
self._dp_sum_query.get_noised_result(
sample_state, self._global_state))
# print(grad_sums)
def normalize(v):
return tf.truediv(
v, tf.cast(self._num_microbatches, tf.float32))
final_grads = tf.nest.map_structure(normalize, grad_sums)
return list(zip(final_grads, var_list))
return DPOptimizerClass
def run(flags_obj):
"""Run ResNet ImageNet training and eval loop using native Keras APIs.
Args:
flags_obj: An object containing parsed flag values.
Raises:
ValueError: If fp16 is passed as it is not currently supported.
NotImplementedError: If some features are not currently supported.
Returns:
Dictionary of training and eval stats.
"""
keras_utils.set_session_config(enable_eager=flags_obj.enable_eager,
enable_xla=flags_obj.enable_xla)
# Execute flag override logic for better model performance
if flags_obj.tf_gpu_thread_mode:
keras_utils.set_gpu_thread_mode_and_count(
per_gpu_thread_count=flags_obj.per_gpu_thread_count,
gpu_thread_mode=flags_obj.tf_gpu_thread_mode,
num_gpus=flags_obj.num_gpus,
datasets_num_private_threads=flags_obj.datasets_num_private_threads
)
common.set_cudnn_batchnorm_mode()
dtype = flags_core.get_tf_dtype(flags_obj)
performance.set_mixed_precision_policy(
flags_core.get_tf_dtype(flags_obj),
flags_core.get_loss_scale(flags_obj, default_for_fp16=128))
data_format = flags_obj.data_format
if data_format is None:
data_format = ('channels_first'
if tf.test.is_built_with_cuda() else 'channels_last')
tf.keras.backend.set_image_data_format(data_format)
# Configures cluster spec for distribution strategy.
_ = distribution_utils.configure_cluster(flags_obj.worker_hosts,
flags_obj.task_index)
strategy = distribution_utils.get_distribution_strategy(
distribution_strategy=flags_obj.distribution_strategy,
num_gpus=flags_obj.num_gpus,
all_reduce_alg=flags_obj.all_reduce_alg,
num_packs=flags_obj.num_packs,
tpu_address=flags_obj.tpu)
if strategy:
# flags_obj.enable_get_next_as_optional controls whether enabling
# get_next_as_optional behavior in DistributedIterator. If true, last
# partial batch can be supported.
strategy.extended.experimental_enable_get_next_as_optional = (
flags_obj.enable_get_next_as_optional)
strategy_scope = distribution_utils.get_strategy_scope(strategy)
# pylint: disable=protected-access
if flags_obj.use_synthetic_data:
distribution_utils.set_up_synthetic_data()
input_fn = common.get_synth_input_fn(
height=imagenet_preprocessing.DEFAULT_IMAGE_SIZE,
width=imagenet_preprocessing.DEFAULT_IMAGE_SIZE,
num_channels=imagenet_preprocessing.NUM_CHANNELS,
num_classes=imagenet_preprocessing.NUM_CLASSES,
dtype=dtype,
drop_remainder=True)
else:
distribution_utils.undo_set_up_synthetic_data()
input_fn = imagenet_preprocessing.input_fn
# When `enable_xla` is True, we always drop the remainder of the batches
# in the dataset, as XLA-GPU doesn't support dynamic shapes.
drop_remainder = flags_obj.enable_xla
# Current resnet_model.resnet50 input format is always channel-last.
# We use keras_application mobilenet model which input format is depends on
# the keras beckend image data format.
# This use_keras_image_data_format flags indicates whether image preprocessor
# output format should be same as the keras backend image data format or just
# channel-last format.
use_keras_image_data_format = (flags_obj.model == 'mobilenet')
train_input_dataset = input_fn(
is_training=True,
data_dir=flags_obj.data_dir,
batch_size=flags_obj.batch_size,
num_epochs=flags_obj.train_epochs,
parse_record_fn=imagenet_preprocessing.get_parse_record_fn(
use_keras_image_data_format=use_keras_image_data_format),
datasets_num_private_threads=flags_obj.datasets_num_private_threads,
dtype=dtype,
drop_remainder=drop_remainder,
tf_data_experimental_slack=flags_obj.tf_data_experimental_slack,
training_dataset_cache=flags_obj.training_dataset_cache,
)
eval_input_dataset = None
if not flags_obj.skip_eval:
eval_input_dataset = input_fn(
is_training=False,
data_dir=flags_obj.data_dir,
batch_size=flags_obj.batch_size,
num_epochs=flags_obj.train_epochs,
parse_record_fn=imagenet_preprocessing.get_parse_record_fn(
use_keras_image_data_format=use_keras_image_data_format),
dtype=dtype,
drop_remainder=drop_remainder)
lr_schedule = common.PiecewiseConstantDecayWithWarmup(
batch_size=flags_obj.batch_size,
epoch_size=imagenet_preprocessing.NUM_IMAGES['train'],
warmup_epochs=common.LR_SCHEDULE[0][1],
boundaries=list(p[1] for p in common.LR_SCHEDULE[1:]),
multipliers=list(p[0] for p in common.LR_SCHEDULE),
compute_lr_on_cpu=True)
steps_per_epoch = (imagenet_preprocessing.NUM_IMAGES['train'] //
flags_obj.batch_size)
with strategy_scope:
if flags_obj.optimizer == 'resnet50_default':
optimizer = common.get_optimizer(lr_schedule)
elif flags_obj.optimizer == 'mobilenet_default':
initial_learning_rate = \
flags_obj.initial_learning_rate_per_sample * flags_obj.batch_size
optimizer = tf.keras.optimizers.SGD(
learning_rate=tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate,
decay_steps=steps_per_epoch *
flags_obj.num_epochs_per_decay,
decay_rate=flags_obj.lr_decay_factor,
staircase=True),
momentum=0.9)
if flags_obj.fp16_implementation == 'graph_rewrite':
# Note: when flags_obj.fp16_implementation == "graph_rewrite", dtype as
# determined by flags_core.get_tf_dtype(flags_obj) would be 'float32'
# which will ensure tf.compat.v2.keras.mixed_precision and
# tf.train.experimental.enable_mixed_precision_graph_rewrite do not double
# up.
optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(
optimizer)
# TODO(hongkuny): Remove trivial model usage and move it to benchmark.
if flags_obj.use_trivial_model:
model = trivial_model.trivial_model(
imagenet_preprocessing.NUM_CLASSES)
elif flags_obj.model == 'resnet50_v1.5':
resnet_model.change_keras_layer(flags_obj.use_tf_keras_layers)
model = resnet_model.resnet50(
num_classes=imagenet_preprocessing.NUM_CLASSES)
elif flags_obj.model == 'mobilenet':
# TODO(kimjaehong): Remove layers attribute when minimum TF version
# support 2.0 layers by default.
model = tf.keras.applications.mobilenet.MobileNet(
weights=None,
classes=imagenet_preprocessing.NUM_CLASSES,
layers=tf.keras.layers)
if flags_obj.pretrained_filepath:
model.load_weights(flags_obj.pretrained_filepath)
if flags_obj.pruning_method == 'polynomial_decay':
if dtype != tf.float32:
raise NotImplementedError(
'Pruning is currently only supported on dtype=tf.float32.')
pruning_params = {
'pruning_schedule':
tfmot.sparsity.keras.PolynomialDecay(
initial_sparsity=flags_obj.pruning_initial_sparsity,
final_sparsity=flags_obj.pruning_final_sparsity,
begin_step=flags_obj.pruning_begin_step,
end_step=flags_obj.pruning_end_step,
frequency=flags_obj.pruning_frequency),
}
model = tfmot.sparsity.keras.prune_low_magnitude(
model, **pruning_params)
elif flags_obj.pruning_method:
raise NotImplementedError(
'Only polynomial_decay is currently supported.')
model.compile(loss='sparse_categorical_crossentropy',
optimizer=optimizer,
metrics=(['sparse_categorical_accuracy']
if flags_obj.report_accuracy_metrics else None),
run_eagerly=flags_obj.run_eagerly)
train_epochs = flags_obj.train_epochs
callbacks = common.get_callbacks(
steps_per_epoch=steps_per_epoch,
pruning_method=flags_obj.pruning_method,
enable_checkpoint_and_export=flags_obj.enable_checkpoint_and_export,
model_dir=flags_obj.model_dir)
# if mutliple epochs, ignore the train_steps flag.
if train_epochs <= 1 and flags_obj.train_steps:
steps_per_epoch = min(flags_obj.train_steps, steps_per_epoch)
train_epochs = 1
num_eval_steps = (imagenet_preprocessing.NUM_IMAGES['validation'] //
flags_obj.batch_size)
validation_data = eval_input_dataset
if flags_obj.skip_eval:
# Only build the training graph. This reduces memory usage introduced by
# control flow ops in layers that have different implementations for
# training and inference (e.g., batch norm).
if flags_obj.set_learning_phase_to_train:
# TODO(haoyuzhang): Understand slowdown of setting learning phase when
# not using distribution strategy.
tf.keras.backend.set_learning_phase(1)
num_eval_steps = None
validation_data = None
if not strategy and flags_obj.explicit_gpu_placement:
# TODO(b/135607227): Add device scope automatically in Keras training loop
# when not using distribition strategy.
no_dist_strat_device = tf.device('/device:GPU:0')
no_dist_strat_device.__enter__()
history = model.fit(train_input_dataset,
epochs=train_epochs,
steps_per_epoch=steps_per_epoch,
callbacks=callbacks,
validation_steps=num_eval_steps,
validation_data=validation_data,
validation_freq=flags_obj.epochs_between_evals,
verbose=2)
eval_output = None
if not flags_obj.skip_eval:
eval_output = model.evaluate(eval_input_dataset,
steps=num_eval_steps,
verbose=2)
if flags_obj.pruning_method:
model = tfmot.sparsity.keras.strip_pruning(model)
if flags_obj.enable_checkpoint_and_export:
if dtype == tf.bfloat16:
logging.warning(
'Keras model.save does not support bfloat16 dtype.')
else:
# Keras model.save assumes a float32 input designature.
export_path = os.path.join(flags_obj.model_dir, 'saved_model')
model.save(export_path, include_optimizer=False)
if not strategy and flags_obj.explicit_gpu_placement:
no_dist_strat_device.__exit__()
stats = common.build_stats(history, eval_output, callbacks)
return stats
def build_convolutional_box_predictor(is_training,
num_classes,
conv_hyperparams_fn,
min_depth,
max_depth,
num_layers_before_predictor,
use_dropout,
dropout_keep_prob,
kernel_size,
box_code_size,
apply_sigmoid_to_scores=False,
class_prediction_bias_init=0.0,
use_depthwise=False,
mask_head_config=None):
"""Builds the ConvolutionalBoxPredictor from the arguments.
Args:
is_training: Indicates whether the BoxPredictor is in training mode.
num_classes: Number of classes.
conv_hyperparams_fn: A function to generate tf-slim arg_scope with
hyperparameters for convolution ops.
min_depth: Minimum feature depth prior to predicting box encodings
and class predictions.
max_depth: Maximum feature depth prior to predicting box encodings
and class predictions. If max_depth is set to 0, no additional
feature map will be inserted before location and class predictions.
num_layers_before_predictor: Number of the additional conv layers before
the predictor.
use_dropout: Option to use dropout or not. Note that a single dropout
op is applied here prior to both box and class predictions, which stands
in contrast to the ConvolutionalBoxPredictor below.
dropout_keep_prob: Keep probability for dropout.
This is only used if use_dropout is True.
kernel_size: Size of final convolution kernel. If the
spatial resolution of the feature map is smaller than the kernel size,
then the kernel size is automatically set to be
min(feature_width, feature_height).
box_code_size: Size of encoding for each box.
apply_sigmoid_to_scores: If True, apply the sigmoid on the output
class_predictions.
class_prediction_bias_init: Constant value to initialize bias of the last
conv2d layer before class prediction.
use_depthwise: Whether to use depthwise convolutions for prediction
steps. Default is False.
mask_head_config: An optional MaskHead object containing configs for mask
head construction.
Returns:
A ConvolutionalBoxPredictor class.
"""
box_prediction_head = box_head.ConvolutionalBoxHead(
is_training=is_training,
box_code_size=box_code_size,
kernel_size=kernel_size,
use_depthwise=use_depthwise)
class_prediction_head = class_head.ConvolutionalClassHead(
is_training=is_training,
num_classes=num_classes,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
kernel_size=kernel_size,
apply_sigmoid_to_scores=apply_sigmoid_to_scores,
class_prediction_bias_init=class_prediction_bias_init,
use_depthwise=use_depthwise)
other_heads = {}
if mask_head_config is not None:
if not mask_head_config.masks_are_class_agnostic:
logging.warning('Note that class specific mask prediction for SSD '
'models is memory consuming.')
other_heads[convolutional_box_predictor.MASK_PREDICTIONS] = (
mask_head.ConvolutionalMaskHead(
is_training=is_training,
num_classes=num_classes,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
kernel_size=kernel_size,
use_depthwise=use_depthwise,
mask_height=mask_head_config.mask_height,
mask_width=mask_head_config.mask_width,
masks_are_class_agnostic=mask_head_config.masks_are_class_agnostic))
return convolutional_box_predictor.ConvolutionalBoxPredictor(
is_training=is_training,
num_classes=num_classes,
box_prediction_head=box_prediction_head,
class_prediction_head=class_prediction_head,
other_heads=other_heads,
conv_hyperparams_fn=conv_hyperparams_fn,
num_layers_before_predictor=num_layers_before_predictor,
min_depth=min_depth,
max_depth=max_depth)
def run(benchmarks: List[Benchmark],
tfx_runner: Optional[tfx_runner_lib.TfxRunner] = None,
pipeline_name: Optional[str] = None,
pipeline_root: Optional[str] = None,
metadata_connection_config: Optional[
metadata_store_pb2.ConnectionConfig] = None,
enable_cache: Optional[bool] = False,
beam_pipeline_args: Optional[List[str]] = None,
**kwargs) -> BenchmarkPipeline:
"""Runs the given benchmarks as part of a single pipeline DAG.
First it concatenates all the benchmark pipelines into a single DAG
benchmark pipeline. Next it executes the workflow via tfx_runner.run().
When the `match` flag is set, matched benchmarks are filtered by name.
When the `runs_per_benchmark` flag is set, each benchmark is run the number
of times specified.
Args:
benchmarks: List of Benchmark instances to include in the suite.
tfx_runner: The TfxRunner instance that defines the platform where
benchmarks are run.
pipeline_name: Name of the benchmark pipeline.
pipeline_root: Path to root directory of the pipeline.
metadata_connection_config: The config to connect to ML metadata.
enable_cache: Whether or not cache is enabled for this run.
beam_pipeline_args: Beam pipeline args for beam jobs within executor.
Executor will use beam DirectRunner as Default.
**kwargs: Additional kwargs forwarded as kwargs to benchmarks.
Returns:
Returns the BenchmarkPipeline that was passed to the tfx_runner.
Raises:
ValueError: If the given tfx_runner is not supported.
"""
if "compile_pipeline" in kwargs:
kwargs.pop("compile_pipeline")
logging.warning(
"The `compile_pipeline` argument DEPRECATED and ignored. "
"Pipelines are now automatically compiled.")
runs_per_benchmark = FLAGS.runs_per_benchmark
if runs_per_benchmark is None:
runs_per_benchmark = int(
os.environ.get("NITROML_RUNS_PER_BENCHMARK", 1))
if not tfx_runner:
logging.info("Setting TFX runner to OSS default: BeamDagRunner.")
tfx_runner = beam_dag_runner.BeamDagRunner()
if runs_per_benchmark <= 0:
raise ValueError(
"runs_per_benchmark must be strictly positive; "
f"got runs_per_benchmark={runs_per_benchmark} instead.")
benchmark_subpipelines = []
for b in benchmarks:
for benchmark_run in range(runs_per_benchmark):
# Call benchmarks with pipeline args.
spec = b(benchmark_run=benchmark_run + 1,
runs_per_benchmark=runs_per_benchmark,
**kwargs)
for benchmark_subpipeline in spec.benchmark_subpipelines:
if re.match(FLAGS.match, benchmark_subpipeline.id):
benchmark_subpipelines.append(benchmark_subpipeline)
if FLAGS.match and not benchmark_subpipelines:
if spec.components_to_always_add:
logging.info(
"No benchmarks matched the pattern '%s'. "
"Running components passed to self.add(..., always=True) only.",
FLAGS.match)
else:
raise ValueError(
f"No benchmarks matched the pattern '{FLAGS.match}'")
benchmark_pipeline = BenchmarkPipeline(
components_to_always_add=spec.components_to_always_add,
benchmark_subpipelines=benchmark_subpipelines,
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
metadata_connection_config=metadata_connection_config,
enable_cache=enable_cache,
beam_pipeline_args=beam_pipeline_args,
**kwargs)
logging.info("NitroML benchmarks:")
for benchmark_name in benchmark_pipeline.benchmark_names:
logging.info("\t%s", benchmark_name)
logging.info("\t\tRUNNING")
dsl_compiler = compiler.Compiler()
pipeline_to_run = dsl_compiler.compile(benchmark_pipeline)
if spec.requested_partial_run:
logging.info("Only running the following nodes:\n%s",
"\n".join(spec.nodes_to_partial_run))
pipeline_to_run = pipeline_filtering.filter_pipeline(
input_pipeline=pipeline_to_run,
pipeline_run_id_fn=(
pipeline_filtering.make_latest_resolver_pipeline_run_id_fn(
benchmark_pipeline.metadata_connection_config)),
skip_nodes=lambda x: x not in set(spec.nodes_to_partial_run))
tfx_runner.run(pipeline_to_run)
return benchmark_pipeline
def __init__(self,
subnetwork_generator,
max_iteration_steps,
logits_dimension=1,
ensemblers=None,
ensemble_strategies=None,
evaluator=None,
adanet_loss_decay=.9,
filepath=None):
"""Initializes an `adanet.keras.Model`.
Args:
subnetwork_generator: The :class:`adanet.subnetwork.Generator` which
defines the candidate subnetworks to train and evaluate at every AdaNet
iteration.
max_iteration_steps: Total number of steps for which to train candidates
per iteration. If :class:`OutOfRange` or :class:`StopIteration` occurs
in the middle, training stops before `max_iteration_steps` steps. When
:code:`None`, it will train the current iteration forever.
logits_dimension: The dimension of the final layer of any subnetworks.
ensemblers: An iterable of :class:`adanet.ensemble.Ensembler` objects that
define how to ensemble a group of subnetworks. If there are multiple,
each should have a different `name` property.
ensemble_strategies: An iterable of :class:`adanet.ensemble.Strategy`
objects that define the candidate ensembles of subnetworks to explore at
each iteration.
evaluator: An :class:`adanet.Evaluator` for candidate selection after all
subnetworks are done training. When :code:`None`, candidate selection
uses a moving average of their :class:`adanet.Ensemble` AdaNet loss
during training instead. In order to use the *AdaNet algorithm* as
described in [Cortes et al., '17], the given :class:`adanet.Evaluator`
must be created with the same dataset partition used during training.
Otherwise, this framework will perform *AdaNet.HoldOut* which uses a
holdout set for candidate selection, but does not benefit from learning
guarantees.
adanet_loss_decay: Float decay for the exponential-moving-average of the
AdaNet objective throughout training. This moving average is a data-
driven way tracking the best candidate with only the training set.
filepath: Directory to save model parameters, graph and etc. This can also
be used to load checkpoints from the directory into a estimator to
continue training a previously saved model.
"""
logging.warning("""The AdaNet Keras API is currently experimental.""")
self._subnetwork_generator = subnetwork_generator
self._max_iteration_steps = max_iteration_steps
self._logits_dimension = logits_dimension
self._ensemblers = ensemblers
self._ensemble_strategies = ensemble_strategies
self._evaluator = evaluator
self._adanet_loss_decay = adanet_loss_decay
self._filepath = filepath
self._model = None
# Use lambdas to defer initialization of Head.
self._loss_head_map = {
"binary_crossentropy":
lambda: tf.estimator.BinaryClassHead,
"mse":
lambda: tf.estimator.RegressionHead(self._logits_dimension),
"mean_squared_error":
lambda: tf.estimator.RegressionHead(self._logits_dimension),
"sparse_categorical_crossentropy":
lambda: tf.estimator.MultiClassHead(self._logits_dimension),
}
def build_convolutional_keras_box_predictor(is_training,
num_classes,
conv_hyperparams,
freeze_batchnorm,
inplace_batchnorm_update,
num_predictions_per_location_list,
min_depth,
max_depth,
num_layers_before_predictor,
use_dropout,
dropout_keep_prob,
kernel_size,
box_code_size,
class_prediction_bias_init=0.0,
use_depthwise=False,
mask_head_config=None,
name='BoxPredictor'):
"""Builds the ConvolutionalBoxPredictor from the arguments.
Args:
is_training: Indicates whether the BoxPredictor is in training mode.
num_classes: Number of classes.
conv_hyperparams: A `hyperparams_builder.KerasLayerHyperparams` object
containing hyperparameters for convolution ops.
freeze_batchnorm: Whether to freeze batch norm parameters during
training or not. When training with a small batch size (e.g. 1), it is
desirable to freeze batch norm update and use pretrained batch norm
params.
inplace_batchnorm_update: Whether to update batch norm moving average
values inplace. When this is false train op must add a control
dependency on tf.graphkeys.UPDATE_OPS collection in order to update
batch norm statistics.
num_predictions_per_location_list: A list of integers representing the
number of box predictions to be made per spatial location for each
feature map.
min_depth: Minimum feature depth prior to predicting box encodings
and class predictions.
max_depth: Maximum feature depth prior to predicting box encodings
and class predictions. If max_depth is set to 0, no additional
feature map will be inserted before location and class predictions.
num_layers_before_predictor: Number of the additional conv layers before
the predictor.
use_dropout: Option to use dropout or not. Note that a single dropout
op is applied here prior to both box and class predictions, which stands
in contrast to the ConvolutionalBoxPredictor below.
dropout_keep_prob: Keep probability for dropout.
This is only used if use_dropout is True.
kernel_size: Size of final convolution kernel. If the
spatial resolution of the feature map is smaller than the kernel size,
then the kernel size is automatically set to be
min(feature_width, feature_height).
box_code_size: Size of encoding for each box.
class_prediction_bias_init: constant value to initialize bias of the last
conv2d layer before class prediction.
use_depthwise: Whether to use depthwise convolutions for prediction
steps. Default is False.
mask_head_config: An optional MaskHead object containing configs for mask
head construction.
name: A string name scope to assign to the box predictor. If `None`, Keras
will auto-generate one from the class name.
Returns:
A ConvolutionalBoxPredictor class.
"""
box_prediction_heads = []
class_prediction_heads = []
mask_prediction_heads = []
other_heads = {}
if mask_head_config is not None:
other_heads[convolutional_box_predictor.MASK_PREDICTIONS] = \
mask_prediction_heads
for stack_index, num_predictions_per_location in enumerate(
num_predictions_per_location_list):
box_prediction_heads.append(
keras_box_head.ConvolutionalBoxHead(
is_training=is_training,
box_code_size=box_code_size,
kernel_size=kernel_size,
conv_hyperparams=conv_hyperparams,
freeze_batchnorm=freeze_batchnorm,
num_predictions_per_location=num_predictions_per_location,
use_depthwise=use_depthwise,
name='ConvolutionalBoxHead_%d' % stack_index))
class_prediction_heads.append(
keras_class_head.ConvolutionalClassHead(
is_training=is_training,
num_classes=num_classes,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
kernel_size=kernel_size,
conv_hyperparams=conv_hyperparams,
freeze_batchnorm=freeze_batchnorm,
num_predictions_per_location=num_predictions_per_location,
class_prediction_bias_init=class_prediction_bias_init,
use_depthwise=use_depthwise,
name='ConvolutionalClassHead_%d' % stack_index))
if mask_head_config is not None:
if not mask_head_config.masks_are_class_agnostic:
logging.warning('Note that class specific mask prediction for SSD '
'models is memory consuming.')
mask_prediction_heads.append(
keras_mask_head.ConvolutionalMaskHead(
is_training=is_training,
num_classes=num_classes,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
kernel_size=kernel_size,
conv_hyperparams=conv_hyperparams,
freeze_batchnorm=freeze_batchnorm,
num_predictions_per_location=num_predictions_per_location,
use_depthwise=use_depthwise,
mask_height=mask_head_config.mask_height,
mask_width=mask_head_config.mask_width,
masks_are_class_agnostic=mask_head_config.
masks_are_class_agnostic,
name='ConvolutionalMaskHead_%d' % stack_index))
return convolutional_keras_box_predictor.ConvolutionalBoxPredictor(
is_training=is_training,
num_classes=num_classes,
box_prediction_heads=box_prediction_heads,
class_prediction_heads=class_prediction_heads,
other_heads=other_heads,
conv_hyperparams=conv_hyperparams,
num_layers_before_predictor=num_layers_before_predictor,
min_depth=min_depth,
max_depth=max_depth,
freeze_batchnorm=freeze_batchnorm,
inplace_batchnorm_update=inplace_batchnorm_update,
name=name)
def build_convolutional_keras_box_predictor(is_training,
num_classes,
conv_hyperparams,
freeze_batchnorm,
inplace_batchnorm_update,
num_predictions_per_location_list,
min_depth,
max_depth,
num_layers_before_predictor,
use_dropout,
dropout_keep_prob,
kernel_size,
box_code_size,
class_prediction_bias_init=0.0,
use_depthwise=False,
mask_head_config=None,
name='BoxPredictor'):
"""Builds the ConvolutionalBoxPredictor from the arguments.
Args:
is_training: Indicates whether the BoxPredictor is in training mode.
num_classes: Number of classes.
conv_hyperparams: A `hyperparams_builder.KerasLayerHyperparams` object
containing hyperparameters for convolution ops.
freeze_batchnorm: Whether to freeze batch norm parameters during
training or not. When training with a small batch size (e.g. 1), it is
desirable to freeze batch norm update and use pretrained batch norm
params.
inplace_batchnorm_update: Whether to update batch norm moving average
values inplace. When this is false train op must add a control
dependency on tf.graphkeys.UPDATE_OPS collection in order to update
batch norm statistics.
num_predictions_per_location_list: A list of integers representing the
number of box predictions to be made per spatial location for each
feature map.
min_depth: Minimum feature depth prior to predicting box encodings
and class predictions.
max_depth: Maximum feature depth prior to predicting box encodings
and class predictions. If max_depth is set to 0, no additional
feature map will be inserted before location and class predictions.
num_layers_before_predictor: Number of the additional conv layers before
the predictor.
use_dropout: Option to use dropout or not. Note that a single dropout
op is applied here prior to both box and class predictions, which stands
in contrast to the ConvolutionalBoxPredictor below.
dropout_keep_prob: Keep probability for dropout.
This is only used if use_dropout is True.
kernel_size: Size of final convolution kernel. If the
spatial resolution of the feature map is smaller than the kernel size,
then the kernel size is automatically set to be
min(feature_width, feature_height).
box_code_size: Size of encoding for each box.
class_prediction_bias_init: constant value to initialize bias of the last
conv2d layer before class prediction.
use_depthwise: Whether to use depthwise convolutions for prediction
steps. Default is False.
mask_head_config: An optional MaskHead object containing configs for mask
head construction.
name: A string name scope to assign to the box predictor. If `None`, Keras
will auto-generate one from the class name.
Returns:
A ConvolutionalBoxPredictor class.
"""
box_prediction_heads = []
class_prediction_heads = []
mask_prediction_heads = []
other_heads = {}
if mask_head_config is not None:
other_heads[convolutional_box_predictor.MASK_PREDICTIONS] = \
mask_prediction_heads
for stack_index, num_predictions_per_location in enumerate(
num_predictions_per_location_list):
box_prediction_heads.append(
keras_box_head.ConvolutionalBoxHead(
is_training=is_training,
box_code_size=box_code_size,
kernel_size=kernel_size,
conv_hyperparams=conv_hyperparams,
freeze_batchnorm=freeze_batchnorm,
num_predictions_per_location=num_predictions_per_location,
use_depthwise=use_depthwise,
name='ConvolutionalBoxHead_%d' % stack_index))
class_prediction_heads.append(
keras_class_head.ConvolutionalClassHead(
is_training=is_training,
num_classes=num_classes,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
kernel_size=kernel_size,
conv_hyperparams=conv_hyperparams,
freeze_batchnorm=freeze_batchnorm,
num_predictions_per_location=num_predictions_per_location,
class_prediction_bias_init=class_prediction_bias_init,
use_depthwise=use_depthwise,
name='ConvolutionalClassHead_%d' % stack_index))
if mask_head_config is not None:
if not mask_head_config.masks_are_class_agnostic:
logging.warning(
'Note that class specific mask prediction for SSD '
'models is memory consuming.')
mask_prediction_heads.append(
keras_mask_head.ConvolutionalMaskHead(
is_training=is_training,
num_classes=num_classes,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
kernel_size=kernel_size,
conv_hyperparams=conv_hyperparams,
freeze_batchnorm=freeze_batchnorm,
num_predictions_per_location=num_predictions_per_location,
use_depthwise=use_depthwise,
mask_height=mask_head_config.mask_height,
mask_width=mask_head_config.mask_width,
masks_are_class_agnostic=mask_head_config.
masks_are_class_agnostic,
name='ConvolutionalMaskHead_%d' % stack_index))
return convolutional_keras_box_predictor.ConvolutionalBoxPredictor(
is_training=is_training,
num_classes=num_classes,
box_prediction_heads=box_prediction_heads,
class_prediction_heads=class_prediction_heads,
other_heads=other_heads,
conv_hyperparams=conv_hyperparams,
num_layers_before_predictor=num_layers_before_predictor,
min_depth=min_depth,
max_depth=max_depth,
freeze_batchnorm=freeze_batchnorm,
inplace_batchnorm_update=inplace_batchnorm_update,
name=name)
def convert_metrics_proto_to_dict(
metrics_for_slice: metrics_for_slice_pb2.MetricsForSlice,
model_name: Optional[Text] = None
) -> Optional[Tuple[slicer.SliceKeyType,
Optional[view_types.MetricsByOutputName]]]:
"""Converts metrics proto to dict."""
model_metrics_map = {}
if metrics_for_slice.metrics:
model_metrics_map[''] = {
'': {
'': _convert_proto_map_to_dict(metrics_for_slice.metrics)
}
}
default_model_name = None
if metrics_for_slice.metric_keys_and_values:
for kv in metrics_for_slice.metric_keys_and_values:
current_model_name = kv.key.model_name
if current_model_name not in model_metrics_map:
model_metrics_map[current_model_name] = {}
output_name = kv.key.output_name
if output_name not in model_metrics_map[current_model_name]:
model_metrics_map[current_model_name][output_name] = {}
sub_key_metrics_map = model_metrics_map[current_model_name][
output_name]
sub_key_id = str(metric_types.SubKey.from_proto(
kv.key.sub_key)) if kv.key.HasField('sub_key') else ''
if sub_key_id not in sub_key_metrics_map:
sub_key_metrics_map[sub_key_id] = {}
if kv.key.is_diff:
if default_model_name is None:
default_model_name = current_model_name
elif default_model_name != current_model_name:
# Setting '' to trigger no match found ValueError below.
default_model_name = ''
metric_name = '{}_diff'.format(kv.key.name)
elif kv.key.HasField('aggregation_type'):
metric_name = '{}_{}'.format(kv.key.aggregation_type,
kv.key.name)
else:
metric_name = kv.key.name
sub_key_metrics_map[sub_key_id][
metric_name] = json_format.MessageToDict(kv.value)
metrics_map = None
keys = list(model_metrics_map.keys())
tmp_model_name = model_name or default_model_name
if tmp_model_name in model_metrics_map:
# Use the provided model name if there is a match.
metrics_map = model_metrics_map[tmp_model_name]
# Add model-independent (e.g. example_count) metrics to all models.
if tmp_model_name and '' in model_metrics_map:
for output_name, output_dict in model_metrics_map[''].items():
for sub_key_id, sub_key_dict in output_dict.items():
for name, value in sub_key_dict.items():
metrics_map.setdefault(output_name,
{}).setdefault(sub_key_id,
{})[name] = value
elif not tmp_model_name and len(keys) == 1:
# Show result of the only model if no model name is specified.
metrics_map = model_metrics_map[keys[0]]
elif keys:
# No match found.
logging.warning(
'Fail to find metrics for model name: %s . '
'Available model names are [%s]', model_name, ', '.join(keys))
return None
return (slicer.deserialize_slice_key(metrics_for_slice.slice_key),
metrics_map)
def create_alerts(config):
""""Creates Stackdriver alerts for logs-based metrics."""
# Stackdriver alerts can't yet be created in Deployment Manager, so create
# them here.
alert_email = config.project.get('stackdriver_alert_email')
if alert_email is None:
logging.warning('No Stackdriver alert email specified, skipping creation '
'of Stackdriver alerts.')
return
project_id = config.project['project_id']
existing_channels_str = runner.run_gcloud_command([
'alpha', 'monitoring', 'channels', 'list', '--format',
'value(name,labels.email_address)'
],
project_id=project_id)
existing_channels = existing_channels_str.split(
'\n') if existing_channels_str else []
email_to_channel = {}
for existing_channel in existing_channels:
channel, email = existing_channel.split()
# assume only one channel exists per email
email_to_channel[email] = channel
if alert_email in email_to_channel:
logging.info('Stackdriver notification channel already exists for %s',
alert_email)
channel = email_to_channel[alert_email]
else:
logging.info('Creating Stackdriver notification channel.')
channel = utils.create_notification_channel(alert_email, project_id)
existing_alerts = runner.run_gcloud_command([
'alpha', 'monitoring', 'policies', 'list', '--format',
'value(displayName)'
],
project_id=project_id).split('\n')
existing_alerts = set(existing_alerts)
logging.info('Creating Stackdriver alerts.')
display_name = 'IAM Policy Change Alert'
if display_name not in existing_alerts:
utils.create_alert_policy(
['global', 'pubsub_topic', 'pubsub_subscription', 'gce_instance'],
'iam-policy-change-count', display_name,
('This policy ensures the designated user/group is notified when IAM '
'policies are altered.'), channel, project_id)
display_name = 'Bucket Permission Change Alert'
if display_name not in existing_alerts:
utils.create_alert_policy(
['gcs_bucket'], 'bucket-permission-change-count', display_name,
('This policy ensures the designated user/group is notified when '
'bucket/object permissions are altered.'), channel, project_id)
display_name = 'Bigquery Update Alert'
if display_name not in existing_alerts:
utils.create_alert_policy(
['global'], 'bigquery-settings-change-count', display_name,
('This policy ensures the designated user/group is notified when '
'Bigquery dataset settings are altered.'), channel, project_id)
for data_bucket in config.project.get('data_buckets', []):
# Every bucket with 'expected_users' has an expected-access alert.
if 'expected_users' not in data_bucket:
continue
bucket_name = get_data_bucket_name(data_bucket, project_id)
metric_name = 'unexpected-access-' + bucket_name
display_name = 'Unexpected Access to {} Alert'.format(bucket_name)
if display_name not in existing_alerts:
utils.create_alert_policy(
['gcs_bucket'], metric_name, display_name,
('This policy ensures the designated user/group is notified when '
'bucket {} is accessed by an unexpected user.'.format(bucket_name)),
channel, project_id)
def _modify_model_output_type(model, inference_output_type=dtypes.float32):
"""Modify model output type."""
if inference_output_type == dtypes.float32:
return
subgraph = model.subgraphs[0]
tensors = subgraph.tensors
operators = subgraph.operators
# Find all dequantize operators
dequant_opcode_idxs = []
for idx, opcode in enumerate(model.operatorCodes):
builtin_code = schema_util.get_builtin_code_from_operator_code(opcode)
if builtin_code == schema_fb.BuiltinOperator.DEQUANTIZE:
dequant_opcode_idxs.append(idx)
if not dequant_opcode_idxs:
raise ValueError("Model output is not dequantized.")
# Validate that the model output is dequantized
output_dequant_ops = []
for op in operators:
# Find operators that dequantize model output
if op.opcodeIndex in dequant_opcode_idxs and \
op.outputs[0] in subgraph.outputs:
# If found, validate that the operator's output type is float
quant_tensor, float_tensor = tensors[op.inputs[0]], tensors[
op.outputs[0]]
float_type = _convert_tflite_enum_type_to_tf_type(
float_tensor.type)
if float_type != dtypes.float32:
raise ValueError(
"Initial model output type must be tf.float32. Expected type for "
"tensor with name '{}' is tf.float32, instead type is {}".
format(float_tensor.name, _get_tf_type_name(float_type)))
# If found, validate that the operator input is quantized and compatible
# with the final model output type
quant_type = _convert_tflite_enum_type_to_tf_type(
quant_tensor.type)
if quant_type not in _MAP_QUANT_TO_IO_TYPES:
raise ValueError(
"Initial model output is not dequantized. Expected type for "
"tensor with name '{}' should be in {}, instead type is {}"
.format(
quant_tensor.name,
tuple(
_get_tf_type_name(t)
for t in _MAP_QUANT_TO_IO_TYPES.keys()),
_get_tf_type_name(quant_type)))
else:
inference_io_types = _MAP_QUANT_TO_IO_TYPES[quant_type]
if inference_output_type not in inference_io_types:
raise ValueError(
"Unsupported `inference_output_type` value. Expected to be in "
"{}, instead got {}.".format(
tuple(
_get_tf_type_name(t)
for t in inference_io_types),
_get_tf_type_name(inference_output_type)))
output_dequant_ops.append(op)
if len(subgraph.outputs) != len(output_dequant_ops):
logging.warning(
"For model outputs containing unsupported operations which cannot be "
"quantized, the `inference_output_type` attribute will default to the "
"original type.")
# Modify model output type
if inference_output_type == dtypes.uint8:
# Find a quantize operator
quant_opcode_idx = -1
for idx, opcode in enumerate(model.operatorCodes):
builtin_code = schema_util.get_builtin_code_from_operator_code(
opcode)
if builtin_code == schema_fb.BuiltinOperator.QUANTIZE:
quant_opcode_idx = idx
break
# Create a quantize operator, if none exist
if quant_opcode_idx == -1:
quant_op = schema_fb.OperatorCodeT()
quant_op.builtinCode = schema_fb.BuiltinOperator.QUANTIZE
quant_op.deprecatedBuiltinCode = schema_fb.BuiltinOperator.QUANTIZE
model.operatorCodes.append(quant_op)
quant_opcode_idx = len(model.operatorCodes) - 1
# Change dequant op (int8 to float) to quant op (int8 to uint8)
for op in output_dequant_ops:
op.opcodeIndex = quant_opcode_idx
int8_quantization = tensors[op.inputs[0]].quantization
uint8_quantization = schema_fb.QuantizationParametersT()
uint8_quantization.scale = [int8_quantization.scale[0]]
uint8_quantization.zeroPoint = [
int8_quantization.zeroPoint[0] + 128
]
tensors[op.outputs[0]].quantization = uint8_quantization
tensors[op.outputs[0]].type = schema_fb.TensorType.UINT8
elif inference_output_type in _MAP_QUANT_TO_IO_TYPES:
# Remove the outputs and the dequant operator
remove_tensors_idxs = set()
for op in output_dequant_ops:
subgraph.outputs[subgraph.outputs == op.outputs[0]] = op.inputs[0]
remove_tensors_idxs.add(op.outputs[0])
operators.remove(op)
# Remove tensors marked for deletion.
_remove_tensors_from_model(model, remove_tensors_idxs)
else:
raise ValueError(
"Unsupported `inference_output_type` value {}.".format(
_get_tf_type_name(inference_output_type)))
def update_eval_config_with_defaults(
eval_config: config_pb2.EvalConfig,
maybe_add_baseline: Optional[bool] = None,
maybe_remove_baseline: Optional[bool] = None,
has_baseline: Optional[bool] = False,
rubber_stamp: Optional[bool] = False) -> config_pb2.EvalConfig:
"""Returns a new config with default settings applied.
a) Add or remove a model_spec according to "has_baseline".
b) Fix the model names (model_spec.name) to tfma.CANDIDATE_KEY and
tfma.BASELINE_KEY.
c) Update the metrics_specs with the fixed model name.
Args:
eval_config: Original eval config.
maybe_add_baseline: DEPRECATED. True to add a baseline ModelSpec to the
config as a copy of the candidate ModelSpec that should already be
present. This is only applied if a single ModelSpec already exists in the
config and that spec doesn't have a name associated with it. When applied
the model specs will use the names tfma.CANDIDATE_KEY and
tfma.BASELINE_KEY. Only one of maybe_add_baseline or maybe_remove_baseline
should be used.
maybe_remove_baseline: DEPRECATED. True to remove a baseline ModelSpec from
the config if it already exists. Removal of the baseline also removes any
change thresholds. Only one of maybe_add_baseline or maybe_remove_baseline
should be used.
has_baseline: True to add a baseline ModelSpec to the config as a copy of
the candidate ModelSpec that should already be present. This is only
applied if a single ModelSpec already exists in the config and that spec
doesn't have a name associated with it. When applied the model specs will
use the names tfma.CANDIDATE_KEY and tfma.BASELINE_KEY. False to remove a
baseline ModelSpec from the config if it already exists. Removal of the
baseline also removes any change thresholds. Only one of has_baseline or
maybe_remove_baseline should be used.
rubber_stamp: True if this model is being rubber stamped. When a model is
rubber stamped diff thresholds will be ignored if an associated baseline
model is not passed.
Raises:
RuntimeError: on missing baseline model for non-rubberstamp cases.
"""
if (not has_baseline and has_change_threshold(eval_config)
and not rubber_stamp):
# TODO(b/173657964): Raise an error instead of logging an error.
raise RuntimeError(
'There are change thresholds, but the baseline is missing. '
'This is allowed only when rubber stamping (first run).')
updated_config = config_pb2.EvalConfig()
updated_config.CopyFrom(eval_config)
# if user requests CIs but doesn't set method, use JACKKNIFE
if (eval_config.options.compute_confidence_intervals.value
and eval_config.options.confidence_intervals.method == config_pb2.
ConfidenceIntervalOptions.UNKNOWN_CONFIDENCE_INTERVAL_METHOD):
updated_config.options.confidence_intervals.method = (
config_pb2.ConfidenceIntervalOptions.JACKKNIFE)
if maybe_add_baseline and maybe_remove_baseline:
raise ValueError(
'only one of maybe_add_baseline and maybe_remove_baseline '
'should be used')
if maybe_add_baseline or maybe_remove_baseline:
logging.warning(
""""maybe_add_baseline" and "maybe_remove_baseline" are deprecated,
please use "has_baseline" instead.""")
if has_baseline:
raise ValueError(
""""maybe_add_baseline" and "maybe_remove_baseline" are ignored if
"has_baseline" is set.""")
if has_baseline is not None:
if has_baseline:
maybe_add_baseline = True
else:
maybe_remove_baseline = True
# Has a baseline model.
if (maybe_add_baseline and len(updated_config.model_specs) == 1
and not updated_config.model_specs[0].name):
baseline = updated_config.model_specs.add()
baseline.CopyFrom(updated_config.model_specs[0])
baseline.name = constants.BASELINE_KEY
baseline.is_baseline = True
updated_config.model_specs[0].name = constants.CANDIDATE_KEY
logging.info(
'Adding default baseline ModelSpec based on the candidate ModelSpec '
'provided. The candidate model will be called "%s" and the baseline '
'will be called "%s": updated_config=\n%s',
constants.CANDIDATE_KEY, constants.BASELINE_KEY, updated_config)
# Does not have a baseline.
if maybe_remove_baseline:
tmp_model_specs = []
for model_spec in updated_config.model_specs:
if not model_spec.is_baseline:
tmp_model_specs.append(model_spec)
del updated_config.model_specs[:]
updated_config.model_specs.extend(tmp_model_specs)
for metrics_spec in updated_config.metrics_specs:
for metric in metrics_spec.metrics:
if metric.threshold.ByteSize():
metric.threshold.ClearField('change_threshold')
for per_slice_threshold in metric.per_slice_thresholds:
if per_slice_threshold.threshold.ByteSize():
per_slice_threshold.threshold.ClearField(
'change_threshold')
for cross_slice_threshold in metric.cross_slice_thresholds:
if cross_slice_threshold.threshold.ByteSize():
cross_slice_threshold.threshold.ClearField(
'change_threshold')
for threshold in metrics_spec.thresholds.values():
if threshold.ByteSize():
threshold.ClearField('change_threshold')
for per_slice_thresholds in metrics_spec.per_slice_thresholds.values(
):
for per_slice_threshold in per_slice_thresholds.thresholds:
if per_slice_threshold.threshold.ByteSize():
per_slice_threshold.threshold.ClearField(
'change_threshold')
for cross_slice_thresholds in metrics_spec.cross_slice_thresholds.values(
):
for cross_slice_threshold in cross_slice_thresholds.thresholds:
if cross_slice_threshold.threshold.ByteSize():
cross_slice_threshold.threshold.ClearField(
'change_threshold')
logging.info(
'Request was made to ignore the baseline ModelSpec and any change '
'thresholds. This is likely because a baseline model was not provided: '
'updated_config=\n%s', updated_config)
if not updated_config.model_specs:
updated_config.model_specs.add()
model_names = []
for spec in updated_config.model_specs:
model_names.append(spec.name)
if len(model_names) == 1 and model_names[0]:
logging.info(
'ModelSpec name "%s" is being ignored and replaced by "" because a '
'single ModelSpec is being used', model_names[0])
updated_config.model_specs[0].name = ''
model_names = ['']
for spec in updated_config.metrics_specs:
if not spec.model_names:
spec.model_names.extend(model_names)
elif len(model_names) == 1:
del spec.model_names[:]
spec.model_names.append('')
return updated_config
def __post_init__(self):
super().__post_init__()
for attr in [
'nums_dense_ftrs', 'nums_sparse_ftrs',
'nums_shallow_tower_sparse_ftrs', 'user_id_column_names',
'shallow_tower_sparse_ftrs_column_names',
'sparse_ftrs_column_names', 'dense_ftrs_column_names',
'user_text_column_names', 'doc_id_column_names',
'doc_text_column_names'
]:
self._set_late_init_attr(attr, [])
# Assemble feature name to num map
self.feature_name2num = dict()
self._set_list_kv_attrs('feature_name2num', 'dense_ftrs_column_names',
'nums_dense_ftrs')
self._set_list_kv_attrs('feature_name2num', 'sparse_ftrs_column_names',
'nums_sparse_ftrs')
self._set_list_kv_attrs('feature_name2num',
'shallow_tower_sparse_ftrs_column_names',
'nums_shallow_tower_sparse_ftrs')
# Assemble feature type to name map
self.feature_type2name = dict()
all_ftr_names = set()
for ftr_type in parsing_utils.get_feature_types():
assert hasattr(
self, ftr_type
), f'{ftr_type} must be defined in DeText argument parser'
ftr_name = getattr(self, ftr_type)
if not ftr_name:
continue
self.feature_type2name[ftr_type] = ftr_name
ftr_name = [ftr_name
] if not isinstance(ftr_name, list) else ftr_name
prev_ftr_num = len(all_ftr_names)
all_ftr_names.update(ftr_name)
if len(all_ftr_names) == prev_ftr_num:
logging.warning(
f"Feature type {ftr_type} has duplicate features with self/other feature types"
)
# Multi-task training: currently only support ranking tasks with both deep and dense features
if self.task_ids:
assert InputFtrType.TASK_ID_COLUMN_NAME in self.feature_type2name, "task_id feature not found for multi-task training"
# Parse task ids an weights from inputs and convert them into a map
task_ids = self.task_ids
raw_weights = self.task_weights if self.task_weights else [
1.0
] * len(task_ids)
task_weights = [
float(wt) / sum(raw_weights) for wt in raw_weights
] # Normalize task weights
# Check size of task_ids and task_weights
assert len(task_ids) == len(
task_weights), "size of task IDs and weights must match"
self.task_weights = task_weights
# Calculate total number of dense features
assert len(self.nums_dense_ftrs) == len(self.dense_ftrs_column_names), \
f'Number of dense features must be consistent in shape and name. E.g., if there are k features in --dense_ftrs, you will also need to specify ' \
f'k integers for --nums_dense_ftrs. Current setting: len(nums_dense_ftrs) = {len(self.nums_dense_ftrs)}, ' \
f'len(dense_ftrs) = {len(self.dense_ftrs_column_names)}'
self.total_num_dense_ftrs = 0
for num_dense_ftrs in self.nums_dense_ftrs:
self.total_num_dense_ftrs += num_dense_ftrs
self.total_num_sparse_ftrs = 0
for num_sparse_ftrs in self.nums_sparse_ftrs:
self.total_num_sparse_ftrs += num_sparse_ftrs
# Get text field numbers
self.has_query = InputFtrType.QUERY_COLUMN_NAME in self.feature_type2name
self.num_doc_fields = len(
self.feature_type2name.get(InputFtrType.DOC_TEXT_COLUMN_NAMES, []))
self.num_user_fields = len(
self.feature_type2name.get(InputFtrType.USER_TEXT_COLUMN_NAMES,
[]))
self.num_text_fields = self.has_query + self.num_doc_fields + self.num_user_fields
if self.ftr_ext != 'bert' and (self.num_doc_fields > 0
or self.num_user_fields > 0
or self.has_query):
assert self.vocab_file or self.vocab_hub_url or self.embedding_hub_url, \
"Must provide vocab or embedding file when text features are provided when not using bert"
# Get id field numbers
self.num_doc_id_fields = len(
self.feature_type2name.get(InputFtrType.DOC_ID_COLUMN_NAMES, []))
self.num_user_id_fields = len(
self.feature_type2name.get(InputFtrType.USER_ID_COLUMN_NAMES, []))
self.num_id_fields = self.num_doc_id_fields + self.num_user_id_fields
if self.num_doc_id_fields > 0 or self.num_user_id_fields > 0:
assert self.vocab_file_for_id_ftr or self.vocab_hub_url_for_id_ftr or self.embedding_hub_url_for_id_ftr, \
"Must provide vocab or embedding file for id features arg when id features are provided"
assert self.use_deep is LateInit, "use_deep is an inferred argument. Do not pass values to this argument"
self.use_deep = any([self.num_text_fields, self.num_id_fields])
# Get indicator of whether dense/sparse features are used
self.use_dense_ftrs = InputFtrType.DENSE_FTRS_COLUMN_NAMES in self.feature_type2name
self.use_sparse_ftrs = InputFtrType.SPARSE_FTRS_COLUMN_NAMES in self.feature_type2name
# Feature normalization
if self.std_file:
# Read normalization file
ftr_mean, ftr_std = parsing_utils.load_ftr_mean_std(self.std_file)
self.ftr_mean = np.array(ftr_mean, dtype=np.float32).tolist()
self.ftr_std = np.array(ftr_std, dtype=np.float32).tolist()
else:
self.ftr_mean = None
self.ftr_std = None
# limitations under the License.
# ==============================================================================
# Only support tensorflow 2.0
# pylint: disable=invalid-name, no-member
r""" entry point for multi-gpu/ multi-machine training """
import sys
import json
import tensorflow as tf
import horovod.tensorflow as hvd
from absl import logging
from athena import HorovodSolver
from athena.main import parse_config, train
if __name__ == "__main__":
logging.set_verbosity(logging.INFO)
if len(sys.argv) < 2:
logging.warning('Usage: python {} config_json_file'.format(sys.argv[0]))
sys.exit()
tf.random.set_seed(1)
json_file = sys.argv[1]
config = None
with open(json_file) as f:
config = json.load(f)
p = parse_config(config)
HorovodSolver.initialize_devices(p.solver_gpu)
#multi-servers training should use hvd.rank()
train(json_file, HorovodSolver, hvd.size(), hvd.rank())
def call_variants(examples_filename,
checkpoint_path,
model,
output_file,
execution_hardware='auto',
batch_size=16,
max_batches=None):
"""Main driver of call_variants."""
# Read a single TFExample to make sure we're not loading an older version.
example_format = tf_utils.get_format_from_examples_path(examples_filename)
if example_format is None:
logging.warning('Unable to read any records from %s. Output will contain '
'zero records.', examples_filename)
io_utils.write_tfrecords([], output_file)
elif example_format != 'raw':
raise ValueError('The TF examples in {} has image/format \'{}\' '
'(expected \'raw\') which means you might need to rerun '
'make_examples to generate the examples again.'.format(
examples_filename, example_format))
if execution_hardware not in _ALLOW_EXECUTION_HARDWARE:
raise ValueError(
'Unexpected execution_hardware={} value. Allowed values are {}'.format(
execution_hardware, ','.join(_ALLOW_EXECUTION_HARDWARE)))
with tf.Graph().as_default():
images, encoded_variants, encoded_alt_allele_indices = prepare_inputs(
examples_filename, model, batch_size, FLAGS.num_readers)
# Create our model and extract the predictions from the model endpoints.
predictions = model.create(images, 3, is_training=False)['Predictions']
# The op for initializing the variables.
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
device_count = {'GPU': 0, 'TPU': 0} if execution_hardware == 'cpu' else {}
config = tf.ConfigProto(device_count=device_count)
with tf.Session(config=config) as sess:
sess.run(init_op)
# Initial the model from the provided checkpoint using our session.
logging.info('Initializing model from %s', checkpoint_path)
model.initialize_from_checkpoint(checkpoint_path, 3, False)(sess)
if execution_hardware == 'accelerator':
if not any(dev.device_type != 'CPU' for dev in sess.list_devices()):
raise ExecutionHardwareError(
'execution_hardware is set to accelerator, but no accelerator '
'was found')
logging.info('Writing calls to %s', output_file)
writer, _ = io_utils.make_proto_writer(output_file)
with writer:
start_time = time.time()
try:
n_batches = 0
n_examples = 0
while max_batches is None or n_batches < max_batches:
n_called = call_batch(sess, writer, encoded_variants,
encoded_alt_allele_indices, predictions)
duration = time.time() - start_time
n_batches += 1
n_examples += n_called
logging.info(
('Processed %s examples in %s batches [%.2f sec per 100]'),
n_examples, n_batches, (100 * duration) / n_examples)
except tf.errors.OutOfRangeError:
logging.info('Done evaluating variants')
def build_weight_shared_convolutional_box_predictor(
is_training,
num_classes,
conv_hyperparams_fn,
depth,
num_layers_before_predictor,
box_code_size,
kernel_size=3,
class_prediction_bias_init=0.0,
use_dropout=False,
dropout_keep_prob=0.8,
share_prediction_tower=False,
apply_batch_norm=True,
use_depthwise=False,
mask_head_config=None,
score_converter_fn=tf.identity,
box_encodings_clip_range=None):
"""Builds and returns a WeightSharedConvolutionalBoxPredictor class.
Args:
is_training: Indicates whether the BoxPredictor is in training mode.
num_classes: number of classes. Note that num_classes *does not*
include the background category, so if groundtruth labels take values
in {0, 1, .., K-1}, num_classes=K (and not K+1, even though the
assigned classification targets can range from {0,... K}).
conv_hyperparams_fn: A function to generate tf-slim arg_scope with
hyperparameters for convolution ops.
depth: depth of conv layers.
num_layers_before_predictor: Number of the additional conv layers before
the predictor.
box_code_size: Size of encoding for each box.
kernel_size: Size of final convolution kernel.
class_prediction_bias_init: constant value to initialize bias of the last
conv2d layer before class prediction.
use_dropout: Whether to apply dropout to class prediction head.
dropout_keep_prob: Probability of keeping activiations.
share_prediction_tower: Whether to share the multi-layer tower between box
prediction and class prediction heads.
apply_batch_norm: Whether to apply batch normalization to conv layers in
this predictor.
use_depthwise: Whether to use depthwise separable conv2d instead of conv2d.
mask_head_config: An optional MaskHead object containing configs for mask
head construction.
score_converter_fn: Callable score converter to perform elementwise op on
class scores.
box_encodings_clip_range: Min and max values for clipping the box_encodings.
Returns:
A WeightSharedConvolutionalBoxPredictor class.
"""
box_prediction_head = box_head.WeightSharedConvolutionalBoxHead(
box_code_size=box_code_size,
kernel_size=kernel_size,
use_depthwise=use_depthwise,
box_encodings_clip_range=box_encodings_clip_range)
class_prediction_head = (class_head.WeightSharedConvolutionalClassHead(
num_classes=num_classes,
kernel_size=kernel_size,
class_prediction_bias_init=class_prediction_bias_init,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
use_depthwise=use_depthwise,
score_converter_fn=score_converter_fn))
other_heads = {}
if mask_head_config is not None:
if not mask_head_config.masks_are_class_agnostic:
logging.warning('Note that class specific mask prediction for SSD '
'models is memory consuming.')
other_heads[convolutional_box_predictor.MASK_PREDICTIONS] = (
mask_head.WeightSharedConvolutionalMaskHead(
num_classes=num_classes,
kernel_size=kernel_size,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
mask_height=mask_head_config.mask_height,
mask_width=mask_head_config.mask_width,
masks_are_class_agnostic=mask_head_config.
masks_are_class_agnostic))
return convolutional_box_predictor.WeightSharedConvolutionalBoxPredictor(
is_training=is_training,
num_classes=num_classes,
box_prediction_head=box_prediction_head,
class_prediction_head=class_prediction_head,
other_heads=other_heads,
conv_hyperparams_fn=conv_hyperparams_fn,
depth=depth,
num_layers_before_predictor=num_layers_before_predictor,
kernel_size=kernel_size,
apply_batch_norm=apply_batch_norm,
share_prediction_tower=share_prediction_tower,
use_depthwise=use_depthwise)
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
if FLAGS.tiny:
yolo = YoloV3Tiny(classes=FLAGS.num_classes)
else:
yolo = YoloV3(classes=FLAGS.num_classes)
yolo.load_weights(FLAGS.weights)
logging.info('weights loaded')
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info('classes loaded')
times = []
try:
vid = cv2.VideoCapture(int(FLAGS.video))
except:
vid = cv2.VideoCapture(FLAGS.video)
out = None
if FLAGS.output:
# by default VideoCapture returns float instead of int
width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(vid.get(cv2.CAP_PROP_FPS))
codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))
fps = 0.0
count = 0
ser = sl.Serial("COM3", 57600)
while True:
_, img = vid.read()
if img is None:
logging.warning("Empty Frame")
time.sleep(0.1)
count+=1
if count < 3:
continue
else:
break
img_in = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_in = tf.expand_dims(img_in, 0)
img_in = transform_images(img_in, FLAGS.size)
t1 = time.time()
boxes, scores, classes, nums = yolo.predict(img_in)
fps = ( fps + (1./(time.time()-t1)) ) / 2
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
for i in range(nums[0]):
#TEST THAT ONLY SENDS DETECTIONS THAT ARE PEOPLE (ANDREW AND JASONS CODE)
if(classes[0][i] == 0 and scores[0][i] >= 0.90):
angles, distances = get_angles_and_distances(img, boxes[0], i)
#IF THEY ARE WITH A 10px SQUARE GO AHEAD AND SHOOT
#IF WE WERE MOVING THE CAMERA/TURRET WE WOULD JUST SEND THE DISTANCES TO THE ARDUINO
#AND LET IT FIGURE OUT WHAT TO DO
if((distances[0] <= 10 and distances[0] >= -10) and (distances[1] <= 10 and distances[1] >= -10)):
ser.write("SHOOT".encode())
print(ser.readline())
print(angles)
print(distances)
print("\n")
img = cv2.putText(img, "FPS: {:.2f}".format(fps), (0, 30),
cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 2)
if FLAGS.output:
out.write(img)
cv2.imshow('output', img)
if cv2.waitKey(1) == ord('q'):
break
cv2.destroyAllWindows()
def build_convolutional_box_predictor(is_training,
num_classes,
conv_hyperparams_fn,
min_depth,
max_depth,
num_layers_before_predictor,
use_dropout,
dropout_keep_prob,
kernel_size,
box_code_size,
apply_sigmoid_to_scores=False,
class_prediction_bias_init=0.0,
use_depthwise=False,
mask_head_config=None):
"""Builds the ConvolutionalBoxPredictor from the arguments.
Args:
is_training: Indicates whether the BoxPredictor is in training mode.
num_classes: Number of classes.
conv_hyperparams_fn: A function to generate tf-slim arg_scope with
hyperparameters for convolution ops.
min_depth: Minimum feature depth prior to predicting box encodings
and class predictions.
max_depth: Maximum feature depth prior to predicting box encodings
and class predictions. If max_depth is set to 0, no additional
feature map will be inserted before location and class predictions.
num_layers_before_predictor: Number of the additional conv layers before
the predictor.
use_dropout: Option to use dropout or not. Note that a single dropout
op is applied here prior to both box and class predictions, which stands
in contrast to the ConvolutionalBoxPredictor below.
dropout_keep_prob: Keep probability for dropout.
This is only used if use_dropout is True.
kernel_size: Size of final convolution kernel. If the
spatial resolution of the feature map is smaller than the kernel size,
then the kernel size is automatically set to be
min(feature_width, feature_height).
box_code_size: Size of encoding for each box.
apply_sigmoid_to_scores: If True, apply the sigmoid on the output
class_predictions.
class_prediction_bias_init: Constant value to initialize bias of the last
conv2d layer before class prediction.
use_depthwise: Whether to use depthwise convolutions for prediction
steps. Default is False.
mask_head_config: An optional MaskHead object containing configs for mask
head construction.
Returns:
A ConvolutionalBoxPredictor class.
"""
box_prediction_head = box_head.ConvolutionalBoxHead(
is_training=is_training,
box_code_size=box_code_size,
kernel_size=kernel_size,
use_depthwise=use_depthwise)
class_prediction_head = class_head.ConvolutionalClassHead(
is_training=is_training,
num_classes=num_classes,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
kernel_size=kernel_size,
apply_sigmoid_to_scores=apply_sigmoid_to_scores,
class_prediction_bias_init=class_prediction_bias_init,
use_depthwise=use_depthwise)
other_heads = {}
if mask_head_config is not None:
if not mask_head_config.masks_are_class_agnostic:
logging.warning('Note that class specific mask prediction for SSD '
'models is memory consuming.')
other_heads[convolutional_box_predictor.MASK_PREDICTIONS] = (
mask_head.ConvolutionalMaskHead(
is_training=is_training,
num_classes=num_classes,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
kernel_size=kernel_size,
use_depthwise=use_depthwise,
mask_height=mask_head_config.mask_height,
mask_width=mask_head_config.mask_width,
masks_are_class_agnostic=mask_head_config.
masks_are_class_agnostic))
return convolutional_box_predictor.ConvolutionalBoxPredictor(
is_training=is_training,
num_classes=num_classes,
box_prediction_head=box_prediction_head,
class_prediction_head=class_prediction_head,
other_heads=other_heads,
conv_hyperparams_fn=conv_hyperparams_fn,
num_layers_before_predictor=num_layers_before_predictor,
min_depth=min_depth,
max_depth=max_depth)
def partition_slices(
metrics: List[metrics_for_slice_pb2.MetricsForSlice],
metric_key: Text,
comparison_type: Text = 'HIGHER',
alpha: float = 0.01,
min_num_examples: int = 1
) -> Tuple[List[SliceComparisonResult], List[SliceComparisonResult]]:
"""Partition slices into significant and non-significant slices.
Args:
metrics: List of slice metrics protos. We assume that the metrics have
MetricValue.confidence_interval field populated. This will be populated
when the metrics computed with confidence intervals enabled.
metric_key: Name of the metric based on which significance testing is done.
comparison_type: Type of comparison indicating if we are looking for slices
whose metric is higher (`HIGHER`) or lower (`LOWER`) than the metric of
the base slice (overall dataset).
alpha: Significance-level for statistical significance testing.
min_num_examples: Minimum number of examples that a slice should have. If it
is set to zero, we don't do any filtering.
Returns:
Tuple containing list of statistically significant and non-significant
slices.
"""
assert comparison_type in ['HIGHER', 'LOWER']
if min_num_examples == 0:
min_num_examples = 1
metrics_dict = {
slicer_lib.deserialize_slice_key(slice_metrics.slice_key): slice_metrics
for slice_metrics in metrics
}
overall_slice_metrics = metrics_dict[()]
del metrics_dict[()]
overall_metrics_dict = _get_metrics_as_dict(overall_slice_metrics)
significant_slices, non_significant_slices = [], []
for slice_key, slice_metrics in metrics_dict.items():
slice_metrics_dict = _get_metrics_as_dict(slice_metrics)
num_examples = int(slice_metrics_dict['example_count'].unsampled_value)
if num_examples < min_num_examples:
continue
# Prune non-interesting slices.
if np.isnan(slice_metrics_dict[metric_key].unsampled_value):
continue
if slice_metrics_dict[metric_key].sample_standard_deviation == 0:
logging.warning('Ignoring slice: %s with standard deviation: %s ',
slice_key,
slice_metrics_dict[metric_key].sample_standard_deviation)
continue
# TODO(pachristopher): Should we use weighted example count?
if slice_metrics_dict['example_count'].unsampled_value <= 1:
logging.warning('Ignoring slice: %s with example count: %s ', slice_key,
slice_metrics_dict['example_count'].unsampled_value)
continue
if overall_metrics_dict[metric_key].sample_standard_deviation == 0:
logging.warning('Overall metric has zero standard deviation.')
# Only consider statistically significant slices.
is_significant, p_value = _is_significant_slice(
slice_metrics_dict[metric_key].unsampled_value,
slice_metrics_dict[metric_key].sample_standard_deviation,
slice_metrics_dict['example_count'].unsampled_value,
overall_metrics_dict[metric_key].unsampled_value,
overall_metrics_dict[metric_key].sample_standard_deviation,
overall_metrics_dict['example_count'].unsampled_value, comparison_type,
alpha)
# Compute effect size for the slice.
effect_size = _compute_effect_size(
slice_metrics_dict[metric_key].unsampled_value,
slice_metrics_dict[metric_key].sample_standard_deviation,
overall_metrics_dict[metric_key].unsampled_value,
overall_metrics_dict[metric_key].sample_standard_deviation)
slice_info = SliceComparisonResult(
slice_key, num_examples, slice_metrics_dict[metric_key].unsampled_value,
overall_metrics_dict[metric_key].unsampled_value, p_value, effect_size,
slice_metrics)
if not is_significant:
non_significant_slices.append(slice_info)
continue
significant_slices.append(slice_info)
return significant_slices, non_significant_slices
def run(iterative_process: adapters.IterativeProcessPythonAdapter,
client_datasets_fn: Callable[[int], List[tf.data.Dataset]],
evaluate_fn: Callable[[Any], Dict[str, float]],
test_fn: Optional[Callable[[Any], Dict[str, float]]] = None):
"""Runs federated training for the given TFF `IterativeProcess` instance.
Args:
iterative_process: An `IterativeProcessPythonAdapter` instance to run.
client_datasets_fn: Function accepting an integer argument (the round
number) and returning a list of client datasets to use as federated data
for that round, and a list of the corresponding client ids.
evaluate_fn: Callable accepting a server state (the `state` of the
`IterationResult`) and returning a dict of evaluation metrics. Used to
compute validation metrics throughout the training process.
test_fn: An optional callable accepting a server state (the `state` of the
`IterationResult`) and returning a dict of test metrics. Used to compute
test metrics at the end of the training process.
Returns:
The `state` of the `IterationResult` representing the result of the training
loop.
"""
if not isinstance(iterative_process,
adapters.IterativeProcessPythonAdapter):
raise TypeError('iterative_process should be type '
'`adapters.IterativeProcessPythonAdapter`.')
if not callable(client_datasets_fn):
raise TypeError('client_datasets_fn should be callable.')
if not callable(evaluate_fn):
raise TypeError('evaluate_fn should be callable.')
if test_fn is not None and not callable(test_fn):
raise TypeError('test_fn should be callable.')
total_rounds = FLAGS.total_rounds
logging.info('Starting iterative_process_training_loop')
initial_state = iterative_process.initialize()
hparam_flags = utils_impl.get_hparam_flags()
hparam_dict = collections.OrderedDict([(name, FLAGS[name].value)
for name in hparam_flags])
checkpoint_mngr, metrics_mngr, summary_writer, profiler = _setup_outputs(
FLAGS.root_output_dir, FLAGS.experiment_name, hparam_dict)
logging.info('Asking checkpoint manager to load checkpoint.')
state, round_num = checkpoint_mngr.load_latest_checkpoint(initial_state)
if state is None:
logging.info('Initializing experiment from scratch.')
state = initial_state
round_num = 0
metrics_mngr.clear_all_rounds()
else:
logging.info('Restarted from checkpoint round %d', round_num)
round_num += 1 # Increment to avoid overwriting current checkpoint
metrics_mngr.clear_rounds_after(last_valid_round_num=round_num - 1)
loop_start_time = time.time()
while round_num < total_rounds:
data_prep_start_time = time.time()
federated_train_data = client_datasets_fn(round_num)
train_metrics = {
'prepare_datasets_secs': time.time() - data_prep_start_time
}
training_start_time = time.time()
prev_model = state.model
# TODO(b/145604851): This try/except is used to circumvent ambiguous TF
# errors during training, and should be removed once the root cause is
# determined (and possibly fixed).
try:
with profiler(round_num):
iteration_result = iterative_process.next(
state, federated_train_data)
except (tf.errors.FailedPreconditionError, tf.errors.NotFoundError,
tf.errors.InternalError) as e:
logging.warning(
'Caught %s exception while running round %d:\n\t%s', type(e),
round_num, e)
continue # restart the loop without incrementing the round number
state = iteration_result.state
round_metrics = iteration_result.metrics
train_metrics['training_secs'] = time.time() - training_start_time
train_metrics['model_delta_l2_norm'] = _compute_numpy_l2_difference(
state.model, prev_model)
train_metrics.update(round_metrics)
# TODO(b/148576550): Wire in client training time metrics into custom
# training loops.
train_metrics.pop('keras_training_time_client_sum_sec')
logging.info('Round {:2d}, {:.2f}s per round in average.'.format(
round_num, (time.time() - loop_start_time) / (round_num + 1)))
if (round_num % FLAGS.rounds_per_checkpoint == 0
or round_num == total_rounds - 1):
save_checkpoint_start_time = time.time()
checkpoint_mngr.save_checkpoint(state, round_num)
train_metrics['save_checkpoint_secs'] = (
time.time() - save_checkpoint_start_time)
metrics = {
'train': train_metrics,
}
if round_num % FLAGS.rounds_per_eval == 0:
evaluate_start_time = time.time()
eval_metrics = evaluate_fn(state)
eval_metrics['evaluate_secs'] = time.time() - evaluate_start_time
metrics['eval'] = eval_metrics
_write_metrics(metrics_mngr, summary_writer, metrics, round_num)
round_num += 1
if test_fn:
test_start_time = time.time()
test_metrics = test_fn(state)
test_metrics['test_evaluate_secs'] = time.time() - test_start_time
metrics = {'test': test_metrics}
_write_metrics(metrics_mngr, summary_writer, metrics, total_rounds)
return state
def remove_duplicate_called_graphs(comp):
"""Deduplicates called graphs for a subset of TFF AST constructs.
Args:
comp: Instance of `building_blocks.ComputationBuildingBlock` whose called
graphs we wish to deduplicate, according to `tree_analysis.trees_equal`.
For `comp` to be eligible here, it must be either a lambda itself whose
body contains no lambdas or blocks, or another computation containing no
lambdas or blocks. This restriction is necessary because
`remove_duplicate_called_graphs` makes no effort to ensure that it is not
pulling references out of their defining scope, except for the case where
`comp` is a lambda itself. This function exits early and logs a warning if
this assumption is violated. Additionally, `comp` must contain only
computations which can be represented in TensorFlow, IE, satisfy the type
restriction in `type_utils.is_tensorflow_compatible_type`.
Returns:
Either a called instance of `building_blocks.CompiledComputation` or a
`building_blocks.CompiledComputation` itself, depending on whether `comp`
is of non-functional or functional type respectively. Additionally, returns
a boolean to match the `transformation_utils.TransformSpec` pattern.
"""
py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
tree_analysis.check_has_unique_names(comp)
name_generator = building_block_factory.unique_name_generator(comp)
if isinstance(comp, building_blocks.Lambda):
comp_to_check = comp.result
else:
comp_to_check = comp
if tree_analysis.count_types(
comp_to_check,
(building_blocks.Lambda, building_blocks.Block)) > 0:
logging.warning(
'The preprocessors have failed to remove called lambdas '
'and blocks; falling back to less efficient, but '
'guaranteed, TensorFlow generation with computation %s.', comp)
return comp, False
leaf_called_graphs = []
def _pack_called_graphs_into_block(inner_comp):
"""Packs deduplicated bindings to called graphs in `leaf_called_graphs`."""
if (isinstance(inner_comp, building_blocks.Call) and isinstance(
inner_comp.function, building_blocks.CompiledComputation)):
for (name, x) in leaf_called_graphs:
if tree_analysis.trees_equal(x, inner_comp):
return building_blocks.Reference(
name, inner_comp.type_signature), True
new_name = next(name_generator)
leaf_called_graphs.append((new_name, inner_comp))
return building_blocks.Reference(new_name,
inner_comp.type_signature), True
return inner_comp, False
if isinstance(comp, building_blocks.Lambda):
transformed_result, _ = transformation_utils.transform_postorder(
comp.result, _pack_called_graphs_into_block)
packed_into_block = building_blocks.Block(leaf_called_graphs,
transformed_result)
parsed, _ = create_tensorflow_representing_block(packed_into_block)
tff_func = building_blocks.Lambda(comp.parameter_name,
comp.parameter_type, parsed)
tf_parser_callable = tree_to_cc_transformations.TFParser()
comp, _ = tree_transformations.insert_called_tf_identity_at_leaves(
tff_func)
tf_generated, _ = transformation_utils.transform_postorder(
comp, tf_parser_callable)
else:
transformed_result, _ = transformation_utils.transform_postorder(
comp, _pack_called_graphs_into_block)
packed_into_block = building_blocks.Block(leaf_called_graphs,
transformed_result)
tf_generated, _ = create_tensorflow_representing_block(
packed_into_block)
return tf_generated, True
def calculate_emt(R: Array, box: Array, **kwargs) -> Array:
"""Calculate the elastic modulus tensor.
energy_fn(R) corresponds to the state around which we are expanding
Args:
R: array of shape (N,dimension) of particle positions. This does not
generalize to arbitrary dimensions and is only implemented for
dimension == 2
dimension == 3
box: A box specifying the shape of the simulation volume. Used to infer
the volume of the unit cell.
Return: C or the tuple (C,converged)
where C is the Elastic modulus tensor as an array of shape (dimension,
dimension,dimension,dimension) that respects the major and minor
symmetries, and converged is a boolean flag (see above).
"""
if not (R.shape[-1] == 2 or R.shape[-1] == 3):
raise AssertionError('Only implemented for 2d and 3d systems.')
if R.dtype is not jnp.dtype('float64'):
logging.warning('Elastic modulus calculations can sometimes lose '
'precision when not using 64-bit precision.')
dim = R.shape[-1]
def setup_energy_fn_general(strain_tensor):
I = jnp.eye(dim, dtype=R.dtype)
@jit
def energy_fn_general(R, gamma):
perturbation = I + gamma * strain_tensor
return energy_fn(R, perturbation=perturbation, **kwargs)
return energy_fn_general
def get_affine_response(strain_tensor):
energy_fn_general = setup_energy_fn_general(strain_tensor)
d2U_dRdgamma = jacfwd(jacrev(energy_fn_general, argnums=0),
argnums=1)(R, 0.)
d2U_dgamma2 = jacfwd(jacrev(energy_fn_general, argnums=1),
argnums=1)(R, 0.)
return d2U_dRdgamma, d2U_dgamma2
strain_tensors = _get_strain_tensor_list(dim, R.dtype)
d2U_dRdgamma_all, d2U_dgamma2_all = vmap(get_affine_response)(
strain_tensors)
#Solve the system of equations.
energy_fn_Ronly = partial(energy_fn, **kwargs)
def hvp(f, primals, tangents):
return jvp(grad(f), primals, tangents)[1]
def hvp_specific_with_tether(v):
return hvp(energy_fn_Ronly, (R, ), (v, )) + tether_strength * v
non_affine_response_all = jsp.sparse.linalg.cg(
vmap(hvp_specific_with_tether), d2U_dRdgamma_all, tol=cg_tol)[0]
#The above line should be functionally equivalent to:
#H0=hessian(energy_fn)(R, box=box, **kwargs).reshape(R.size,R.size) \
# + tether_strength * jnp.identity(R.size)
#non_affine_response_all = jnp.transpose(jnp.linalg.solve(
# H0,
# jnp.transpose(d2U_dRdgamma_all))
# )
residual = jnp.linalg.norm(
vmap(hvp_specific_with_tether)(non_affine_response_all) -
d2U_dRdgamma_all)
converged = residual / jnp.linalg.norm(d2U_dRdgamma_all) < cg_tol
response_all = d2U_dgamma2_all - jnp.einsum(
"nij,nij->n", d2U_dRdgamma_all, non_affine_response_all)
vol_0 = quantity.volume(dim, box)
response_all = response_all / vol_0
C = _convert_responses_to_elastic_constants(response_all)
# JAX does not allow proper runtime error handling in jitted function.
# Instead, if the user requests a gradient check and the check fails,
# we convert C into jnp.nan's. While this doesn't raise an exception,
# it at least is very "loud".
if gradient_check is not None:
maxgrad = jnp.amax(jnp.abs(grad(energy_fn)(R, **kwargs)))
C = lax.cond(maxgrad > gradient_check, lambda _: jnp.nan * C,
lambda _: C, None)
if check_convergence:
return C, converged
else:
return C
def build_weight_shared_convolutional_box_predictor(
is_training,
num_classes,
conv_hyperparams_fn,
depth,
num_layers_before_predictor,
box_code_size,
kernel_size=3,
class_prediction_bias_init=0.0,
use_dropout=False,
dropout_keep_prob=0.8,
share_prediction_tower=False,
apply_batch_norm=True,
use_depthwise=False,
mask_head_config=None,
score_converter_fn=tf.identity,
box_encodings_clip_range=None):
"""Builds and returns a WeightSharedConvolutionalBoxPredictor class.
Args:
is_training: Indicates whether the BoxPredictor is in training mode.
num_classes: number of classes. Note that num_classes *does not*
include the background category, so if groundtruth labels take values
in {0, 1, .., K-1}, num_classes=K (and not K+1, even though the
assigned classification targets can range from {0,... K}).
conv_hyperparams_fn: A function to generate tf-slim arg_scope with
hyperparameters for convolution ops.
depth: depth of conv layers.
num_layers_before_predictor: Number of the additional conv layers before
the predictor.
box_code_size: Size of encoding for each box.
kernel_size: Size of final convolution kernel.
class_prediction_bias_init: constant value to initialize bias of the last
conv2d layer before class prediction.
use_dropout: Whether to apply dropout to class prediction head.
dropout_keep_prob: Probability of keeping activiations.
share_prediction_tower: Whether to share the multi-layer tower between box
prediction and class prediction heads.
apply_batch_norm: Whether to apply batch normalization to conv layers in
this predictor.
use_depthwise: Whether to use depthwise separable conv2d instead of conv2d.
mask_head_config: An optional MaskHead object containing configs for mask
head construction.
score_converter_fn: Callable score converter to perform elementwise op on
class scores.
box_encodings_clip_range: Min and max values for clipping the box_encodings.
Returns:
A WeightSharedConvolutionalBoxPredictor class.
"""
box_prediction_head = box_head.WeightSharedConvolutionalBoxHead(
box_code_size=box_code_size,
kernel_size=kernel_size,
use_depthwise=use_depthwise,
box_encodings_clip_range=box_encodings_clip_range)
class_prediction_head = (
class_head.WeightSharedConvolutionalClassHead(
num_classes=num_classes,
kernel_size=kernel_size,
class_prediction_bias_init=class_prediction_bias_init,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
use_depthwise=use_depthwise,
score_converter_fn=score_converter_fn))
other_heads = {}
if mask_head_config is not None:
if not mask_head_config.masks_are_class_agnostic:
logging.warning('Note that class specific mask prediction for SSD '
'models is memory consuming.')
other_heads[convolutional_box_predictor.MASK_PREDICTIONS] = (
mask_head.WeightSharedConvolutionalMaskHead(
num_classes=num_classes,
kernel_size=kernel_size,
use_dropout=use_dropout,
dropout_keep_prob=dropout_keep_prob,
mask_height=mask_head_config.mask_height,
mask_width=mask_head_config.mask_width,
masks_are_class_agnostic=mask_head_config.masks_are_class_agnostic))
return convolutional_box_predictor.WeightSharedConvolutionalBoxPredictor(
is_training=is_training,
num_classes=num_classes,
box_prediction_head=box_prediction_head,
class_prediction_head=class_prediction_head,
other_heads=other_heads,
conv_hyperparams_fn=conv_hyperparams_fn,
depth=depth,
num_layers_before_predictor=num_layers_before_predictor,
kernel_size=kernel_size,
apply_batch_norm=apply_batch_norm,
share_prediction_tower=share_prediction_tower,
use_depthwise=use_depthwise)
def main(_argv):
id_img = 0
physical_devices = tf.config.experimental.list_physical_devices('GPU')
for physical_device in physical_devices:
tf.config.experimental.set_memory_growth(physical_device, True)
if FLAGS.tiny:
yolo = YoloV3Tiny(classes=FLAGS.num_classes)
else:
yolo = YoloV3(classes=FLAGS.num_classes)
yolo.load_weights(FLAGS.weights)
logging.info('weights loaded')
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info('classes loaded')
logging.info('initialization connexion at {}:{}'.format(
FLAGS.ip, str(FLAGS.port)))
connexion_1 = time.time()
sk = nt.init_connexion(FLAGS.ip, FLAGS.port)
sk.close()
connexion_2 = time.time()
connexion = connexion_2 - connexion_1
logging.info("connected to {} port {} in {:.3f}ms".format(
FLAGS.ip, FLAGS.port, connexion))
times = []
try:
vid = cv2.VideoCapture(int(FLAGS.video))
except:
vid = cv2.VideoCapture(FLAGS.video)
out = None
if FLAGS.output:
# by default VideoCapture returns float instead of int
width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(vid.get(cv2.CAP_PROP_FPS))
codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))
while True:
id_img = id_img + 1
_, img = vid.read()
if img is None:
logging.warning("Empty Frame")
time.sleep(0.1)
continue
img_in = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_in = tf.expand_dims(img_in, 0)
img_in = transform_images(img_in, FLAGS.size)
t1 = time.time()
boxes, scores, classes, nums = yolo.predict(img_in)
t2 = time.time()
process = ThreadSending(img, (boxes, scores, classes, nums),
class_names, id_img)
process.setDaemon(True)
process.start()
times.append(t2 - t1)
times = times[-20:]
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
img = cv2.putText(
img, "Time prim: {:.2f}ms".format(sum(times) / len(times) * 1000),
(0, 30), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (255, 255, 255), 2)
if FLAGS.output:
out.write(img)
cv2.imshow('output', img)
if cv2.waitKey(1) == ord('q'):
break
cv2.destroyAllWindows()
# width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
# height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
# fps = int(vid.get(cv2.CAP_PROP_FPS))
# codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
# out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))
fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
out = cv2.VideoWriter(
'/home/duanyajun/文档/目标识别项目/自己改写代码/mobilenet_yolov1/ball_water/109.mp4',
fourcc, int(vid.get(cv2.CAP_PROP_FPS)),
(int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)),
int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))))
while True:
_, img = vid.read()
start = time.time()
if img is None:
logging.warning("Empty Frame")
time.sleep(0.1)
continue
img = image_video(model, img)
end = time.time()
print("time: {:.03f}s, fps: {:.03f}".format(end - start,
1 / (end - start)))
#if FLAGS.output:
# out.write(img)
out.write(img)
#cv2.imshow('output', img)
#if cv2.waitKey(1) == ord('q'):
#break
cv2.destroyAllWindows()
def yield_predictions_from_estimator(predictions, total, log_interval=10):
"""Yields predictions from Estimator.predict, with added error correction.
This function handles the case of Estimator.predict occasionally restarting,
causing results to be yielded out of order.
Args:
predictions: the return value of Estimator.predict. An iterable of dicts.
Each dict MUST have a 'result_idx' attribute, used to track result order.
total (int): total expected number of elements to yield from predictions.
log_interval: log every this many seconds.
Yields:
the same dicts yielded from Estimator.predict, but in the right order. The
result_idx element is removed from every dict.
"""
predictions_iter = iter(predictions)
total_yielded = 0
start_time = time.time()
last_log_timestamp = time.time()
while total_yielded < total:
try:
result = next(predictions_iter)
except StopIteration:
raise ValueError(
'Estimator.predict terminated before we got all results.')
result_idx = result.pop('result_idx')
if result_idx == total_yielded:
# If results are always emitted from Estimator.predict in the same
# order that they were fed into the Estimator, then we should always
# expect result_idx to equal total_yielded. However, this does not always
# happen, so we handle that in the `else` case below.
yield result
total_yielded += 1
# Log progress.
current_time = time.time()
if current_time - last_log_timestamp > log_interval:
total_time = current_time - start_time
log_msg = 'Yielded {} results in {:.2f} secs.'.format(
total_yielded, total_time)
logging.info(log_msg)
last_log_timestamp = current_time
else:
# If results start to arrive out of order, something has gone wrong.
if result_idx < total_yielded:
# This can happen if the TPU worker dies, causing Estimator.predict to
# restart from the beginning. In this case, we just don't yield
# anything on this step. Instead, we keep pulling things from the
# iterator until we are back to where we were.
if result_idx == 0:
logging.warning('TPU worker seems to have restarted.')
elif result_idx > total_yielded:
# Something has gone really wrong.
raise ValueError(
'Estimator.predict has somehow missed a result.')
def __init__(self,
cell,
inference_batch_size=1,
mode=modes.Modes.TRAINING,
pad_time_dim=None,
state_shape=None,
ring_buffer_size_in_time_dim=None,
use_one_step=True,
state_name_tag='ExternalState',
**kwargs):
super(Stream, self).__init__(**kwargs)
self.cell = cell
self.inference_batch_size = inference_batch_size
self.mode = mode
self.pad_time_dim = pad_time_dim
self.state_shape = state_shape
self.ring_buffer_size_in_time_dim = ring_buffer_size_in_time_dim
self.use_one_step = use_one_step
self.state_name_tag = state_name_tag
if not use_one_step and isinstance(
self.cell,
(tf.keras.layers.Flatten, tf.keras.layers.GlobalMaxPooling2D,
tf.keras.layers.GlobalAveragePooling2D)):
raise ValueError(
'Flatten, GlobalMaxPooling2D, GlobalAveragePooling2D '
'can be used only with use_one_step = True '
'because they are executed one time per inference call '
'and produce only one output in time dim, whereas conv '
'can produce multiple outputs in time dim, '
'so conv can be used with use_one_step = False or True')
if self.ring_buffer_size_in_time_dim is not None:
# it is a special case when ring_buffer_size_in_time_dim is specified
# outside of the layer in this case we just build a ring buffer
# and do not check what is the type of the cell
pass
elif isinstance(
cell,
(tf.keras.layers.Conv1D, tf.keras.layers.Conv2D,
tf.keras.layers.DepthwiseConv2D, tf.keras.layers.SeparableConv2D,
tf.keras.layers.SeparableConv1D,
average_pooling2d.AveragePooling2D)):
padding = cell.get_config()['padding']
strides = cell.get_config()['strides']
if self.mode not in (modes.Modes.TRAINING,
modes.Modes.NON_STREAM_INFERENCE):
if padding != 'valid':
raise ValueError('conv/cell padding has to be valid,'
'padding has to be set by pad_time_dim')
if self.use_one_step:
if strides[0] > 1:
raise ValueError(
'Stride in time dim greater than 1 '
'in streaming mode with use_one_step=True'
' is not supported, set use_one_step=False')
dilation_rate = cell.get_config()['dilation_rate']
kernel_size = cell.get_config()['kernel_size']
if self.use_one_step:
# effective kernel size in time dimension
self.ring_buffer_size_in_time_dim = dilation_rate[0] * (
kernel_size[0] - 1) + 1
else:
# Streaming of strided or 1 step conv.
# Assuming input length is a multiple of strides (otherwise streaming
# conv is not meaningful), setting to this value (instead of
# dilation_rate[0] * (kernel_size[0] - 1)) ensures that we do not
# ignore the `strides - 1` rightmost (and hence most recent) valid
# input samples.
self.ring_buffer_size_in_time_dim = max(
0,
dilation_rate[0] * (kernel_size[0] - 1) - (strides[0] - 1))
elif isinstance(self.cell, tf.keras.layers.AveragePooling2D):
strides = cell.get_config()['strides']
pool_size = cell.get_config()['pool_size']
if self.mode not in (modes.Modes.TRAINING,
modes.Modes.NON_STREAM_INFERENCE
) and strides[0] != pool_size[0]:
raise ValueError(
'Stride in time %d must = pool size in time %d' %
(strides[0], pool_size[0]))
# effective kernel size in time dimension
self.ring_buffer_size_in_time_dim = pool_size[0]
elif isinstance(
self.cell,
(tf.keras.layers.Flatten, tf.keras.layers.GlobalMaxPooling2D,
tf.keras.layers.GlobalAveragePooling2D)):
# effective kernel size in time dimension
if self.state_shape:
self.ring_buffer_size_in_time_dim = self.state_shape[1]
else:
raise ValueError('Cell is not supported ', cell)
if self.ring_buffer_size_in_time_dim == 1:
logging.warning(
'There is no need to use Stream on time dim with size 1')
def compute_return_and_advantage(self, next_time_steps, value_preds):
"""Compute the Monte Carlo return and advantage.
Normalazation will be applied to the computed returns and advantages if
it's enabled.
Args:
next_time_steps: batched tensor of TimeStep tuples after action is taken.
value_preds: Batched value prediction tensor. Should have one more entry
in time index than time_steps, with the final value corresponding to the
value prediction of the final state.
Returns:
tuple of (return, normalized_advantage), both are batched tensors.
"""
discounts = next_time_steps.discount * tf.constant(
self._discount_factor, dtype=tf.float32)
rewards = next_time_steps.reward
if self._debug_summaries:
# Summarize rewards before they get normalized below.
tf.compat.v2.summary.histogram(
name='rewards', data=rewards, step=self.train_step_counter)
tf.compat.v2.summary.scalar(
name='rewards_mean',
data=tf.reduce_mean(rewards),
step=self.train_step_counter)
# Normalize rewards if self._reward_normalizer is defined.
if self._reward_normalizer:
rewards = self._reward_normalizer.normalize(
rewards, center_mean=False, clip_value=self._reward_norm_clipping)
if self._debug_summaries:
tf.compat.v2.summary.histogram(
name='rewards_normalized',
data=rewards,
step=self.train_step_counter)
tf.compat.v2.summary.scalar(
name='rewards_normalized_mean',
data=tf.reduce_mean(rewards),
step=self.train_step_counter)
# Make discount 0.0 at end of each episode to restart cumulative sum
# end of each episode.
episode_mask = common.get_episode_mask(next_time_steps)
discounts *= episode_mask
# Compute Monte Carlo returns. Data from incomplete trajectories, not
# containing the end of an episode will also be used, with a bootstrapped
# estimation from the last value.
# Note that when a trajectory driver is used, then the final step is
# terminal, the bootstrapped estimation will not be used, as it will be
# multiplied by zero (the discount on the last step).
final_value_bootstrapped = value_preds[:, -1]
returns = value_ops.discounted_return(
rewards,
discounts,
time_major=False,
final_value=final_value_bootstrapped)
if self._debug_summaries:
tf.compat.v2.summary.histogram(
name='returns', data=returns, step=self.train_step_counter)
# Compute advantages.
advantages = self.compute_advantages(rewards, returns, discounts,
value_preds)
normalized_advantages = _normalize_advantages(advantages, axes=(0, 1))
if self._debug_summaries:
tf.compat.v2.summary.histogram(
name='advantages', data=advantages, step=self.train_step_counter)
tf.compat.v2.summary.histogram(
name='advantages_normalized',
data=normalized_advantages,
step=self.train_step_counter)
# Return TD-Lambda returns if both use_td_lambda_return and use_gae.
if self._use_td_lambda_return:
if not self._use_gae:
logging.warning('use_td_lambda_return was True, but use_gae was '
'False. Using Monte Carlo return.')
else:
returns = tf.add(
advantages, value_preds[:, :-1], name='td_lambda_returns')
return returns, normalized_advantages
def restore_ckpt(model,
ckpt_path_or_file,
ema_decay=0.9998,
skip_mismatch=True):
"""Restore variables from a given checkpoint.
Args:
model: the keras model to be restored.
ckpt_path_or_file: the path or file for checkpoint.
ema_decay: ema decay rate. If None or zero or negative value, disable ema.
skip_mismatch: whether to skip variables if shape mismatch.
"""
if ckpt_path_or_file == '_':
logging.info('Running test: do not load any ckpt.')
return
if tf.io.gfile.isdir(ckpt_path_or_file):
ckpt_path_or_file = tf.train.latest_checkpoint(ckpt_path_or_file)
var_shape_map = tf.train.load_checkpoint(
ckpt_path_or_file).get_variable_to_shape_map()
if '_CHECKPOINTABLE_OBJECT_GRAPH' in var_shape_map:
model.load_weights(ckpt_path_or_file)
else:
if ema_decay > 0:
ema = tf.train.ExponentialMovingAverage(decay=0.0)
ema_vars = get_ema_vars(model)
var_dict = {
average_name(ema, var): var
for (ref, var) in ema_vars.items()
}
else:
ema_vars = get_ema_vars(model)
var_dict = {
var.name.split(':')[0]: var
for (ref, var) in ema_vars.items()
}
# add variables that not in var_dict
for v in model.weights:
if v.ref() not in ema_vars:
var_dict[v.name.split(':')[0]] = v
# try to load graph-based checkpoint with ema support,
# else load checkpoint via keras.load_weights which doesn't support ema.
for i, (key, var) in enumerate(var_dict.items()):
if key in var_shape_map:
if var_shape_map[key] != var.shape:
msg = 'Shape mismatch: %s' % key
if skip_mismatch:
logging.warning(msg)
else:
raise ValueError(msg)
else:
var.assign(tf.train.load_variable(ckpt_path_or_file, key))
if i < 10:
logging.info('Init %s from %s (%s)', var.name, key,
ckpt_path_or_file)
else:
msg = 'Not found %s in %s' % (key, ckpt_path_or_file)
if skip_mismatch:
logging.warning(msg)
else:
raise KeyError(msg)
def setUpClass(cls):
super().setUpClass()
cls.compiled_modules = {}
if cls._modules_to_compile:
for name, (ctor, exported_names,
backends) in cls._modules_to_compile.items():
# Setup the debug directory.
debug_parent_dir = FLAGS.debug_dir
if not debug_parent_dir:
debug_parent_dir = FLAGS.test_tmpdir
debug_parent_dir = os.path.join(debug_parent_dir, cls.__name__)
try:
os.makedirs(debug_parent_dir)
except IOError:
logging.exception(
"Error creating crash reproducer dir for: %s",
debug_parent_dir)
# Setup crash reproducer and global debug dir.
crash_reproducer_path = os.path.join(debug_parent_dir,
name + "_reproducer.mlir")
compiler.Context.default_crash_reproducer_path = crash_reproducer_path
global global_debug_dir
global_debug_dir = debug_parent_dir
try:
# Compile.
# Expand backend names to BackendInfo objects.
def _resolve(backend_spec):
if isinstance(backend_spec, BackendInfo):
return backend_spec
# Handle the string form.
return BackendInfo.ALL[backend_spec]
override_backends = get_override_backends()
if override_backends is not None:
backends = override_backends
elif backends is None:
backends = list(BackendInfo.ALL.keys())
backends = [_resolve(backend) for backend in backends]
# if "tf" is specified as a only backend then
# we will test it always against "tf" by adding "tf_also".
if len(backends) == 1 and "tf" == backends[0].name:
backends.append(BackendInfo.ALL["tf_also"])
available_backends = get_available_backends()
backends = [
backend for backend in backends
if backend in available_backends
]
if not backends:
# If no backends are available, then to avoid errors down the line,
# just use "tf", which should always be safe.
backends = [BackendInfo.ALL["tf"]]
logging.warning(
"Falling back to just `tf` backend because no other requested backends are available. Available backends '%s'",
[backend.name for backend in available_backends])
cls.compiled_modules[name] = dict([
(backend.name,
CompiledModule.create(ctor, exported_names, backend))
for backend in backends
])
finally:
# Disable crash reproducer (to avoid inadvertently overwriting this
# path on a subsequent interaction).
compiler.Context.default_crash_reproducer_path = None
global_debug_dir = None
def run_ncf(_):
"""Run NCF training and eval with Keras."""
# TODO(seemuch): Support different train and eval batch sizes
if FLAGS.eval_batch_size != FLAGS.batch_size:
logging.warning(
"The Keras implementation of NCF currently does not support batch_size "
"!= eval_batch_size ({} vs. {}). Overriding eval_batch_size to match "
"batch_size".format(FLAGS.eval_batch_size, FLAGS.batch_size)
)
FLAGS.eval_batch_size = FLAGS.batch_size
params = ncf_common.parse_flags(FLAGS)
batch_size = params["batch_size"]
# ncf_common rounds eval_batch_size (this is needed due to a reshape during
# eval). This carries over that rounding to batch_size as well.
params['batch_size'] = params['eval_batch_size']
num_users, num_items, num_train_steps, num_eval_steps, producer = (
ncf_common.get_inputs(params))
params["num_users"], params["num_items"] = num_users, num_items
producer.start()
model_helpers.apply_clean(flags.FLAGS)
batches_per_step = params["batches_per_step"]
train_input_dataset, eval_input_dataset = _get_train_and_eval_data(producer,
params)
# It is required that for distributed training, the dataset must call
# batch(). The parameter of batch() here is the number of replicas involed,
# such that each replica evenly gets a slice of data.
train_input_dataset = train_input_dataset.batch(batches_per_step)
eval_input_dataset = eval_input_dataset.batch(batches_per_step)
strategy = ncf_common.get_distribution_strategy(params)
with distribution_utils.get_strategy_scope(strategy):
keras_model = _get_keras_model(params)
optimizer = tf.keras.optimizers.Adam(
learning_rate=params["learning_rate"],
beta_1=params["beta1"],
beta_2=params["beta2"],
epsilon=params["epsilon"])
time_callback = keras_utils.TimeHistory(batch_size, FLAGS.log_steps)
keras_model.compile(
loss=_keras_loss,
metrics=[_get_metric_fn(params)],
optimizer=optimizer)
history = keras_model.fit(train_input_dataset,
epochs=FLAGS.train_epochs,
callbacks=[
IncrementEpochCallback(producer),
time_callback],
verbose=2)
logging.info("Training done. Start evaluating")
eval_results = keras_model.evaluate(
eval_input_dataset,
steps=num_eval_steps,
verbose=2)
logging.info("Keras evaluation is done.")
stats = build_stats(history, eval_results, time_callback)
return stats
def __init__(self,
head,
subnetwork_generator,
max_iteration_steps,
ensemblers=None,
ensemble_strategies=None,
evaluator=None,
report_materializer=None,
metric_fn=None,
force_grow=False,
replicate_ensemble_in_training=False,
adanet_loss_decay=.9,
model_dir=None,
report_dir=None,
config=None,
use_tpu=True,
eval_on_tpu=True,
export_to_tpu=True,
train_batch_size=None,
eval_batch_size=None,
predict_batch_size=None,
embedding_config_spec=None,
debug=False,
enable_ensemble_summaries=True,
enable_subnetwork_summaries=True,
export_subnetwork_logits=False,
export_subnetwork_last_layer=True,
global_step_combiner_fn=tf.math.reduce_mean,
max_iterations=None,
replay_config=None,
**kwargs):
self._use_tpu = use_tpu
if not self._use_tpu:
logging.warning(
"This adanet.TPUEstimator is meant to be used for running on TPU. "
"If you want to run on CPU/GPU, use adanet.Estimator instead.")
# TPUEstimator modifies config under the hood. We keep track of it here so
# we can use it from _create_temp_run_config.
self._original_config = config or tf_compat.v1.estimator.tpu.RunConfig()
self._eval_on_tpu = eval_on_tpu if self._use_tpu else False
self._export_to_tpu = export_to_tpu
self._train_batch_size = train_batch_size or 0
self._eval_batch_size = eval_batch_size or train_batch_size or 0
self._predict_batch_size = (
predict_batch_size or eval_batch_size or train_batch_size or 0)
self._embedding_config_spec = embedding_config_spec
if self._embedding_config_spec:
logging.warning(
"TPU does not support inference with TPUEmbedding. Force setting "
"`export_to_tpu=False` so no TPU SavedModel will be exported.")
self._export_to_tpu = False
from tensorflow_estimator.python.estimator.tpu import tpu_estimator # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
super(TPUEstimator, self).__init__(
head=head,
subnetwork_generator=subnetwork_generator,
max_iteration_steps=max_iteration_steps,
ensemblers=ensemblers,
ensemble_strategies=ensemble_strategies,
evaluator=evaluator,
report_materializer=report_materializer,
metric_fn=metric_fn,
force_grow=force_grow,
replicate_ensemble_in_training=replicate_ensemble_in_training,
adanet_loss_decay=adanet_loss_decay,
model_dir=model_dir,
report_dir=report_dir,
config=self._original_config,
use_tpu=self._use_tpu,
eval_on_tpu=self._eval_on_tpu,
export_to_tpu=self._export_to_tpu,
export_saved_model_api_version=(
tpu_estimator.ExportSavedModelApiVersion.V2),
train_batch_size=self._train_batch_size,
eval_batch_size=self._eval_batch_size,
predict_batch_size=self._predict_batch_size,
embedding_config_spec=self._embedding_config_spec,
debug=debug,
enable_ensemble_summaries=enable_ensemble_summaries,
enable_subnetwork_summaries=enable_subnetwork_summaries,
export_subnetwork_logits=export_subnetwork_logits,
export_subnetwork_last_layer=export_subnetwork_last_layer,
global_step_combiner_fn=global_step_combiner_fn,
max_iterations=max_iterations,
replay_config=replay_config,
**kwargs)
def __init__(self,
policy: tf_policy.TFPolicy,
batch_size: Optional[int] = None,
use_nest_path_signatures: bool = True,
seed: Optional[types.Seed] = None,
train_step: Optional[tf.Variable] = None,
input_fn_and_spec: Optional[InputFnAndSpecType] = None,
metadata: Optional[Dict[Text, tf.Variable]] = None):
"""Initialize PolicySaver for TF policy `policy`.
Args:
policy: A TF Policy.
batch_size: The number of batch entries the policy will process at a time.
This must be either `None` (unknown batch size) or a python integer.
use_nest_path_signatures: SavedModel spec signatures will be created based
on the sructure of the specs. Otherwise all specs must have unique
names.
seed: Random seed for the `policy.action` call, if any (this should
usually be `None`, except for testing).
train_step: Variable holding the train step for the policy. The value
saved will be set at the time `saver.save` is called. If not provided,
train_step defaults to -1. Note since the train step must be a variable
it is not safe to create it directly in TF1 so in that case this is a
required parameter.
input_fn_and_spec: A `(input_fn, tensor_spec)` tuple where input_fn is a
function that takes inputs according to tensor_spec and converts them to
the `(time_step, policy_state)` tuple that is used as the input to the
action_fn. When `input_fn_and_spec` is set, `tensor_spec` is the input
for the action signature. When `input_fn_and_spec is None`, the action
signature takes as input `(time_step, policy_state)`.
metadata: A dictionary of `tf.Variables` to be saved along with the
policy.
Raises:
TypeError: If `policy` is not an instance of TFPolicy.
TypeError: If `metadata` is not a dictionary of tf.Variables.
ValueError: If use_nest_path_signatures is not used and any of the
following `policy` specs are missing names, or the names collide:
`policy.time_step_spec`, `policy.action_spec`,
`policy.policy_state_spec`, `policy.info_spec`.
ValueError: If `batch_size` is not either `None` or a python integer > 0.
"""
if not isinstance(policy, tf_policy.TFPolicy):
raise TypeError('policy is not a TFPolicy. Saw: %s' %
type(policy))
if (batch_size is not None
and (not isinstance(batch_size, int) or batch_size < 1)):
raise ValueError(
'Expected batch_size == None or python int > 0, saw: %s' %
(batch_size, ))
action_fn_input_spec = (policy.time_step_spec,
policy.policy_state_spec)
if use_nest_path_signatures:
action_fn_input_spec = _rename_spec_with_nest_paths(
action_fn_input_spec)
else:
_check_spec(action_fn_input_spec)
# Make a shallow copy as we'll be making some changes in-place.
saved_policy = tf.Module()
saved_policy.collect_data_spec = copy.copy(policy.collect_data_spec)
saved_policy.policy_state_spec = copy.copy(policy.policy_state_spec)
if train_step is None:
if not common.has_eager_been_enabled():
raise ValueError('train_step is required in TF1 and must be a '
'`tf.Variable`: %s' % train_step)
train_step = tf.Variable(
-1,
trainable=False,
dtype=tf.int64,
aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
shape=())
elif not isinstance(train_step, tf.Variable):
raise ValueError('train_step must be a TensorFlow variable: %s' %
train_step)
# We will need the train step for the Checkpoint object.
self._train_step = train_step
saved_policy.train_step = self._train_step
self._metadata = metadata or {}
for key, value in self._metadata.items():
if not isinstance(key, str):
raise TypeError('Keys of metadata must be strings: %s' % key)
if not isinstance(value, tf.Variable):
raise TypeError('Values of metadata must be tf.Variable: %s' %
value)
saved_policy.metadata = self._metadata
if batch_size is None:
get_initial_state_fn = policy.get_initial_state
get_initial_state_input_specs = (tf.TensorSpec(
dtype=tf.int32, shape=(), name='batch_size'), )
else:
get_initial_state_fn = functools.partial(policy.get_initial_state,
batch_size=batch_size)
get_initial_state_input_specs = ()
get_initial_state_fn = common.function()(get_initial_state_fn)
original_action_fn = policy.action
if seed is not None:
def action_fn(time_step, policy_state):
time_step = cast(ts.TimeStep, time_step)
return original_action_fn(time_step, policy_state, seed=seed)
else:
action_fn = original_action_fn
def distribution_fn(time_step, policy_state):
"""Wrapper for policy.distribution() in the SavedModel."""
try:
time_step = cast(ts.TimeStep, time_step)
outs = policy.distribution(time_step=time_step,
policy_state=policy_state)
return tf.nest.map_structure(_composite_distribution, outs)
except (TypeError, NotImplementedError) as e:
# TODO(b/156526399): Move this to just the policy.distribution() call
# once tfp.experimental.as_composite() properly handles LinearOperator*
# components as well as TransformedDistributions.
logging.warning(
'WARNING: Could not serialize policy.distribution() for policy '
'"%s". Calling saved_model.distribution() will raise the following '
'assertion error: %s', policy, e)
@common.function()
def _raise():
tf.Assert(False, [str(e)])
return ()
outs = _raise()
# We call get_concrete_function() for its side effect: to ensure the proper
# ConcreteFunction is stored in the SavedModel.
get_initial_state_fn.get_concrete_function(
*get_initial_state_input_specs)
train_step_fn = common.function(
lambda: saved_policy.train_step).get_concrete_function()
get_metadata_fn = common.function(
lambda: saved_policy.metadata).get_concrete_function()
batched_time_step_spec = tf.nest.map_structure(
lambda spec: add_batch_dim(spec, [batch_size]),
policy.time_step_spec)
batched_time_step_spec = cast(ts.TimeStep, batched_time_step_spec)
batched_policy_state_spec = tf.nest.map_structure(
lambda spec: add_batch_dim(spec, [batch_size]),
policy.policy_state_spec)
policy_step_spec = policy.policy_step_spec
policy_state_spec = policy.policy_state_spec
if use_nest_path_signatures:
batched_time_step_spec = _rename_spec_with_nest_paths(
batched_time_step_spec)
batched_policy_state_spec = _rename_spec_with_nest_paths(
batched_policy_state_spec)
policy_step_spec = _rename_spec_with_nest_paths(policy_step_spec)
policy_state_spec = _rename_spec_with_nest_paths(policy_state_spec)
else:
_check_spec(batched_time_step_spec)
_check_spec(batched_policy_state_spec)
_check_spec(policy_step_spec)
_check_spec(policy_state_spec)
if input_fn_and_spec is not None:
# Store a signature based on input_fn_and_spec
@common.function()
def polymorphic_action_fn(example):
action_inputs = input_fn_and_spec[0](example)
tf.nest.map_structure(
lambda spec, t: tf.Assert(spec.is_compatible_with(t[
0]), [t]), action_fn_input_spec, action_inputs)
return action_fn(*action_inputs)
@common.function()
def polymorphic_distribution_fn(example):
action_inputs = input_fn_and_spec[0](example)
tf.nest.map_structure(
lambda spec, t: tf.Assert(spec.is_compatible_with(t[
0]), [t]), action_fn_input_spec, action_inputs)
return distribution_fn(*action_inputs)
batched_input_spec = tf.nest.map_structure(
lambda spec: add_batch_dim(spec, [batch_size]),
input_fn_and_spec[1])
# We call get_concrete_function() for its side effect: to ensure the
# proper ConcreteFunction is stored in the SavedModel.
polymorphic_action_fn.get_concrete_function(
example=batched_input_spec)
polymorphic_distribution_fn.get_concrete_function(
example=batched_input_spec)
action_input_spec = (input_fn_and_spec[1], )
else:
action_input_spec = action_fn_input_spec
if batched_policy_state_spec:
# Store the signature with a required policy state spec
polymorphic_action_fn = common.function()(action_fn)
polymorphic_action_fn.get_concrete_function(
time_step=batched_time_step_spec,
policy_state=batched_policy_state_spec)
polymorphic_distribution_fn = common.function()(
distribution_fn)
polymorphic_distribution_fn.get_concrete_function(
time_step=batched_time_step_spec,
policy_state=batched_policy_state_spec)
else:
# Create a polymorphic action_fn which you can call as
# restored.action(time_step)
# or
# restored.action(time_step, ())
# (without retracing the inner action twice)
@common.function()
def polymorphic_action_fn(
time_step, policy_state=batched_policy_state_spec):
return action_fn(time_step, policy_state)
polymorphic_action_fn.get_concrete_function(
time_step=batched_time_step_spec,
policy_state=batched_policy_state_spec)
polymorphic_action_fn.get_concrete_function(
time_step=batched_time_step_spec)
@common.function()
def polymorphic_distribution_fn(
time_step, policy_state=batched_policy_state_spec):
return distribution_fn(time_step, policy_state)
polymorphic_distribution_fn.get_concrete_function(
time_step=batched_time_step_spec,
policy_state=batched_policy_state_spec)
polymorphic_distribution_fn.get_concrete_function(
time_step=batched_time_step_spec)
signatures = {
# CompositeTensors aren't well supported by old-style signature
# mechanisms, so we do not have a signature for policy.distribution.
'action':
_function_with_flat_signature(polymorphic_action_fn,
input_specs=action_input_spec,
output_spec=policy_step_spec,
include_batch_dimension=True,
batch_size=batch_size),
'get_initial_state':
_function_with_flat_signature(
get_initial_state_fn,
input_specs=get_initial_state_input_specs,
output_spec=policy_state_spec,
include_batch_dimension=False),
'get_train_step':
_function_with_flat_signature(train_step_fn,
input_specs=(),
output_spec=train_step.dtype,
include_batch_dimension=False),
'get_metadata':
_function_with_flat_signature(get_metadata_fn,
input_specs=(),
output_spec=tf.nest.map_structure(
lambda v: v.dtype,
self._metadata),
include_batch_dimension=False),
}
saved_policy.action = polymorphic_action_fn
saved_policy.distribution = polymorphic_distribution_fn
saved_policy.get_initial_state = get_initial_state_fn
saved_policy.get_train_step = train_step_fn
saved_policy.get_metadata = get_metadata_fn
# Adding variables as an attribute to facilitate updating them.
saved_policy.model_variables = policy.variables()
# TODO(b/156779400): Move to a public API for accessing all trackable leaf
# objects (once it's available). For now, we have no other way of tracking
# objects like Tables, Vocabulary files, etc.
try:
saved_policy._all_assets = policy._unconditional_checkpoint_dependencies # pylint: disable=protected-access
except AttributeError as e:
if '_self_unconditional' in str(e):
logging.warning(
'Unable to capture all trackable objects in policy "%s". This '
'may be okay. Error: %s', policy, e)
else:
raise e
self._policy = saved_policy
self._signatures = signatures
self._action_input_spec = action_input_spec
self._policy_step_spec = policy_step_spec
self._policy_state_spec = policy_state_spec
from uncertainty_baselines.datasets.places import Places365Dataset
from uncertainty_baselines.datasets.random import RandomGaussianImageDataset
from uncertainty_baselines.datasets.random import RandomRademacherImageDataset
from uncertainty_baselines.datasets.svhn import SvhnDataset
from uncertainty_baselines.datasets.test_utils import DatasetTest
from uncertainty_baselines.datasets.toxic_comments import CivilCommentsDataset
from uncertainty_baselines.datasets.toxic_comments import CivilCommentsIdentitiesDataset
from uncertainty_baselines.datasets.toxic_comments import WikipediaToxicityDataset
# pylint: enable=g-bad-import-order
try:
from uncertainty_baselines.datasets.smcalflow import MultiWoZDataset # pylint: disable=g-import-not-at-top
from uncertainty_baselines.datasets.smcalflow import SMCalflowDataset # pylint: disable=g-import-not-at-top
except ImportError:
logging.warning(
'Skipped importing the SMCalflow dataset due to ImportError. Try '
'installing uncertainty baselines with the `datasets` extras.',
exc_info=True)
try:
# Try to import datasets depending on librosa.
from uncertainty_baselines.datasets.speech_commands import SpeechCommandsDataset # pylint: disable=g-import-not-at-top
except ImportError:
logging.warning(
'Skipped importing the Speech Commands dataset due to ImportError. Try '
'installing uncertainty baselines with the `datasets` extras.',
exc_info=True)
except OSError:
logging.warning(
'Skipped importing the Speech Commands dataset due to OSError.',
exc_info=True)
def __init__(
self,
episode_length,
modality_types,
confidence_threshold,
output_size,
worlds,
targets,
compute_distance,
should_draw_detections,
dataset_root,
labelmap_path,
reward_collision,
reward_goal_range,
num_detection_classes,
segmentation_file_name,
detection_folder_name,
actions,
targets_file_name,
eval_init_points_file_name=None,
shaped_reward=False,
):
"""Instantiates the environment for ActiveVision Dataset.
Args:
episode_length: the length of each episode.
modality_types: a list of the strings where each entry indicates the name
of the modalities to be loaded. Valid entries are "sseg", "det",
"depth", "image", "distance", and "prev_action". "distance" should be
used for computing metrics in tf agents.
confidence_threshold: Consider detections more than confidence_threshold
for potential targets.
output_size: Resolution of the output image.
worlds: List of the name of the worlds.
targets: List of the target names. Each entry is a string label of the
target category (e.g. 'fridge', 'microwave', so on).
compute_distance: If True, outputs the distance of the view to the goal.
should_draw_detections (bool): If True, the image returned for the
observation will contains the bounding boxes.
dataset_root: the path to the root folder of the dataset.
labelmap_path: path to the dictionary that converts label strings to
indexes.
reward_collision: the reward the agents get after hitting an obstacle.
It should be a non-positive number.
reward_goal_range: the number of steps from goal, such that the agent is
considered to have reached the goal. If the agent's distance is less
than the specified goal range, the episode is also finishes by setting
done = True.
num_detection_classes: number of classes that detector outputs.
segmentation_file_name: the name of the file that contains the semantic
information. The file should be in the dataset_root/Meta/ folder.
detection_folder_name: Name of the folder that contains the detections
for each world. The folder should be under dataset_root/Meta/ folder.
actions: The list of the action names. Valid entries are listed in
SUPPORTED_ACTIONS.
targets_file_name: the name of the file that contains the annotated
targets. The file should be in the dataset_root/Meta/Folder
eval_init_points_file_name: The name of the file that contains the initial
points for evaluating the performance of the agent. If set to None,
episodes start at random locations. Should be only set for evaluation.
shaped_reward: Whether to add delta goal distance to the reward each step.
Raises:
ValueError: If one of the targets are not available in the annotated
targets or the modality names are not from the domain specified above.
ValueError: If one of the actions is not in SUPPORTED_ACTIONS.
ValueError: If the reward_collision is a positive number.
ValueError: If there is no action other than stop provided.
"""
if reward_collision > 0:
raise ValueError('"reward" for collision should be non positive')
if reward_goal_range < 0:
logging.warning('environment does not terminate the episode if the agent '
'is too close to the environment')
if not modality_types:
raise ValueError('modality names can not be empty')
for name in modality_types:
if name not in SUPPORTED_MODALITIES:
raise ValueError('invalid modality type: {}'.format(name))
actions_other_than_stop_found = False
for a in actions:
if a != 'stop':
actions_other_than_stop_found = True
if a not in SUPPORTED_ACTIONS:
raise ValueError('invalid action %s', a)
if not actions_other_than_stop_found:
raise ValueError('environment needs to have actions other than stop.')
super(ActiveVisionDatasetEnv, self).__init__()
self._episode_length = episode_length
self._modality_types = set(modality_types)
self._confidence_threshold = confidence_threshold
self._output_size = output_size
self._dataset_root = dataset_root
self._worlds = worlds
self._targets = targets
self._all_graph = {}
for world in self._worlds:
with tf.gfile.Open(_get_json_path(self._dataset_root, world), 'r') as f:
file_content = f.read()
file_content = file_content.replace('.jpg', '')
io = StringIO(file_content)
self._all_graph[world] = json.load(io)
self._cur_world = ''
self._cur_image_id = ''
self._cur_graph = None # Loaded by _update_graph
self._steps_taken = 0
self._last_action_success = True
self._category_index = _init_category_index(labelmap_path)
self._category_map = dict(
[(c, i) for i, c in enumerate(self._category_index)])
self._detection_cache = {}
if not ActiveVisionDatasetEnv.cached_data:
ActiveVisionDatasetEnv.cached_data = read_cached_data(
True, self._dataset_root, segmentation_file_name, targets_file_name,
self._output_size)
cached_data = ActiveVisionDatasetEnv.cached_data
self._world_id_dict = cached_data['world_id_dict']
self._depth_images = cached_data[task_env.ModalityTypes.DEPTH]
self._semantic_segmentations = cached_data[
task_env.ModalityTypes.SEMANTIC_SEGMENTATION]
self._annotated_targets = cached_data['targets']
self._cached_imgs = cached_data[task_env.ModalityTypes.IMAGE]
self._graph_cache = {}
self._compute_distance = compute_distance
self._should_draw_detections = should_draw_detections
self._reward_collision = reward_collision
self._reward_goal_range = reward_goal_range
self._num_detection_classes = num_detection_classes
self._actions = actions
self._detection_folder_name = detection_folder_name
self._shaped_reward = shaped_reward
self._eval_init_points = None
if eval_init_points_file_name is not None:
self._eval_init_index = 0
init_points_path = os.path.join(self._dataset_root, 'Meta',
eval_init_points_file_name + '.npy')
with tf.gfile.Open(init_points_path) as points_file:
data = np.load(points_file).item()
self._eval_init_points = []
for world in self._worlds:
for goal in self._targets:
if world in self._annotated_targets[goal]:
for image_id in data[world]:
self._eval_init_points.append((world, image_id[0], goal))
logging.info('loaded %d eval init points', len(self._eval_init_points))
self.action_space = gym.spaces.Discrete(len(self._actions))
obs_shapes = {}
if task_env.ModalityTypes.SEMANTIC_SEGMENTATION in self._modality_types:
obs_shapes[task_env.ModalityTypes.SEMANTIC_SEGMENTATION] = gym.spaces.Box(
low=0, high=255, shape=(self._output_size, self._output_size, 1))
if task_env.ModalityTypes.OBJECT_DETECTION in self._modality_types:
obs_shapes[task_env.ModalityTypes.OBJECT_DETECTION] = gym.spaces.Box(
low=0,
high=255,
shape=(self._output_size, self._output_size,
self._num_detection_classes))
if task_env.ModalityTypes.DEPTH in self._modality_types:
obs_shapes[task_env.ModalityTypes.DEPTH] = gym.spaces.Box(
low=0,
high=_MAX_DEPTH_VALUE,
shape=(self._output_size, self._output_size, 2))
if task_env.ModalityTypes.IMAGE in self._modality_types:
obs_shapes[task_env.ModalityTypes.IMAGE] = gym.spaces.Box(
low=0, high=255, shape=(self._output_size, self._output_size, 3))
if task_env.ModalityTypes.GOAL in self._modality_types:
obs_shapes[task_env.ModalityTypes.GOAL] = gym.spaces.Box(
low=0, high=1., shape=(len(self._targets),))
if task_env.ModalityTypes.PREV_ACTION in self._modality_types:
obs_shapes[task_env.ModalityTypes.PREV_ACTION] = gym.spaces.Box(
low=0, high=1., shape=(len(self._actions) + 1,))
if task_env.ModalityTypes.DISTANCE in self._modality_types:
obs_shapes[task_env.ModalityTypes.DISTANCE] = gym.spaces.Box(
low=0, high=255, shape=(1,))
self.observation_space = gym.spaces.Dict(obs_shapes)
self._prev_action = np.zeros((len(self._actions) + 1), dtype=np.float32)
# Loading all the poses.
all_poses = {}
for world in self._worlds:
all_poses[world] = read_all_poses(self._dataset_root, world)
self._cached_poses = all_poses
self._vertex_to_pose = {}
self._pose_to_vertex = {}
def run_ncf(_):
"""Run NCF training and eval with Keras."""
keras_utils.set_session_config(enable_xla=FLAGS.enable_xla)
if FLAGS.seed is not None:
print("Setting tf seed")
tf.random.set_seed(FLAGS.seed)
# TODO(seemuch): Support different train and eval batch sizes
if FLAGS.eval_batch_size != FLAGS.batch_size:
logging.warning(
"The Keras implementation of NCF currently does not support batch_size "
"!= eval_batch_size ({} vs. {}). Overriding eval_batch_size to match "
"batch_size".format(FLAGS.eval_batch_size, FLAGS.batch_size))
FLAGS.eval_batch_size = FLAGS.batch_size
params = ncf_common.parse_flags(FLAGS)
model_helpers.apply_clean(flags.FLAGS)
strategy = distribution_utils.get_distribution_strategy(
distribution_strategy=FLAGS.distribution_strategy,
num_gpus=FLAGS.num_gpus)
params["distribute_strategy"] = strategy
if not keras_utils.is_v2_0() and strategy is not None:
logging.error(
"NCF Keras only works with distribution strategy in TF 2.0")
return
if (params["keras_use_ctl"]
and (not keras_utils.is_v2_0() or strategy is None)):
logging.error(
"Custom training loop only works with tensorflow 2.0 and dist strat."
)
return
# ncf_common rounds eval_batch_size (this is needed due to a reshape during
# eval). This carries over that rounding to batch_size as well. This is the
# per device batch size
params["batch_size"] = params["eval_batch_size"]
batch_size = params["batch_size"]
time_callback = keras_utils.TimeHistory(batch_size, FLAGS.log_steps)
callbacks = [time_callback]
producer, input_meta_data = None, None
generate_input_online = params["train_dataset_path"] is None
if generate_input_online:
# Start data producing thread.
num_users, num_items, num_train_steps, num_eval_steps, producer = (
ncf_common.get_inputs(params))
producer.start()
per_epoch_callback = IncrementEpochCallback(producer)
callbacks.append(per_epoch_callback)
else:
assert params["eval_dataset_path"] and params["input_meta_data_path"]
with tf.io.gfile.GFile(params["input_meta_data_path"], "rb") as reader:
input_meta_data = json.loads(reader.read().decode("utf-8"))
num_users = input_meta_data["num_users"]
num_items = input_meta_data["num_items"]
params["num_users"], params["num_items"] = num_users, num_items
(train_input_dataset, eval_input_dataset, num_train_steps, num_eval_steps) = \
(ncf_input_pipeline.create_ncf_input_data(
params, producer, input_meta_data))
steps_per_epoch = None if generate_input_online else num_train_steps
if FLAGS.early_stopping:
early_stopping_callback = CustomEarlyStopping(
"val_HR_METRIC", desired_value=FLAGS.hr_threshold)
callbacks.append(early_stopping_callback)
with distribution_utils.get_strategy_scope(strategy):
keras_model = _get_keras_model(params)
optimizer = tf.keras.optimizers.Adam(
learning_rate=params["learning_rate"],
beta_1=params["beta1"],
beta_2=params["beta2"],
epsilon=params["epsilon"])
if params["keras_use_ctl"]:
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
reduction="sum", from_logits=True)
train_input_iterator = strategy.make_dataset_iterator(
train_input_dataset)
eval_input_iterator = strategy.make_dataset_iterator(
eval_input_dataset)
@tf.function
def train_step():
"""Called once per step to train the model."""
def step_fn(features):
"""Computes loss and applied gradient per replica."""
with tf.GradientTape() as tape:
softmax_logits = keras_model(features)
labels = features[rconst.TRAIN_LABEL_KEY]
loss = loss_object(
labels,
softmax_logits,
sample_weight=features[rconst.VALID_POINT_MASK])
loss *= (1.0 /
(batch_size * strategy.num_replicas_in_sync))
grads = tape.gradient(loss, keras_model.trainable_variables)
# Converting gradients to dense form helps in perf on GPU for NCF
grads = neumf_model.sparse_to_dense_grads(
list(zip(grads, keras_model.trainable_variables)))
optimizer.apply_gradients(grads)
return loss
per_replica_losses = strategy.experimental_run(
step_fn, train_input_iterator)
mean_loss = strategy.reduce(tf.distribute.ReduceOp.SUM,
per_replica_losses,
axis=None)
return mean_loss
@tf.function
def eval_step():
"""Called once per eval step to compute eval metrics."""
def step_fn(features):
"""Computes eval metrics per replica."""
softmax_logits = keras_model(features)
in_top_k, metric_weights = metric_fn(
softmax_logits, features[rconst.DUPLICATE_MASK], params)
hr_sum = tf.reduce_sum(in_top_k * metric_weights)
hr_count = tf.reduce_sum(metric_weights)
return hr_sum, hr_count
per_replica_hr_sum, per_replica_hr_count = (
strategy.experimental_run(step_fn, eval_input_iterator))
hr_sum = strategy.reduce(tf.distribute.ReduceOp.SUM,
per_replica_hr_sum,
axis=None)
hr_count = strategy.reduce(tf.distribute.ReduceOp.SUM,
per_replica_hr_count,
axis=None)
return hr_sum, hr_count
time_callback.on_train_begin()
for epoch in range(FLAGS.train_epochs):
for cb in callbacks:
cb.on_epoch_begin(epoch)
# As NCF dataset is sampled with randomness, not repeating
# data elements in each epoch has significant impact on
# convergence. As so, offline-generated TF record files
# contains all epoch worth of data. Thus we do not need
# to initialize dataset when reading from tf record files.
if generate_input_online:
train_input_iterator.initialize()
train_loss = 0
for step in range(num_train_steps):
time_callback.on_batch_begin(step + epoch * num_train_steps)
train_loss += train_step()
time_callback.on_batch_end(step + epoch * num_train_steps)
train_loss /= num_train_steps
logging.info("Done training epoch %s, epoch loss=%s.", epoch + 1,
train_loss)
eval_input_iterator.initialize()
hr_sum = 0
hr_count = 0
for _ in range(num_eval_steps):
step_hr_sum, step_hr_count = eval_step()
hr_sum += step_hr_sum
hr_count += step_hr_count
logging.info("Done eval epoch %s, hr=%s.", epoch + 1,
hr_sum / hr_count)
if (FLAGS.early_stopping
and float(hr_sum / hr_count) > params["hr_threshold"]):
break
time_callback.on_train_end()
eval_results = [None, hr_sum / hr_count]
else:
with distribution_utils.get_strategy_scope(strategy):
# TODO(b/138957587): Remove when force_v2_in_keras_compile is on longer
# a valid arg for this model. Also remove as a valid flag.
if FLAGS.force_v2_in_keras_compile is not None:
keras_model.compile(optimizer=optimizer,
run_eagerly=FLAGS.run_eagerly,
experimental_run_tf_function=FLAGS.
force_v2_in_keras_compile)
else:
keras_model.compile(optimizer=optimizer,
run_eagerly=FLAGS.run_eagerly)
history = keras_model.fit(train_input_dataset,
epochs=FLAGS.train_epochs,
steps_per_epoch=steps_per_epoch,
callbacks=callbacks,
validation_data=eval_input_dataset,
validation_steps=num_eval_steps,
verbose=2)
logging.info("Training done. Start evaluating")
eval_results = keras_model.evaluate(eval_input_dataset,
steps=num_eval_steps,
verbose=2)
logging.info("Keras evaluation is done.")
if history and history.history:
train_history = history.history
train_loss = train_history["loss"][-1]
stats = build_stats(train_loss, eval_results, time_callback)
return stats
def _resolve_overlapping_variants(overlapping_variants):
"""Yields variants with compatible haplotypes, if possible.
Args:
overlapping_variants: list(Variant). A non-empty list of Variant protos in
coordinate-sorted order that overlap on the reference genome and are
predicted to contain alternate allele genotypes.
Yields:
Variant protos in coordinate-sorted order that try to resolve incompatible
haplotypes.
"""
# Short circuit the simplest case: A single variant in a region is compatible
# with itself by definition.
if len(overlapping_variants) == 1:
yield overlapping_variants[0]
return
# If the actual genotype calls are compatible, we can safely return those
# since they would be the most likely configuration also when restricting to
# only valid configurations of genotype calls.
calculator = _VariantCompatibilityCalculator(overlapping_variants)
nonref_counts = [_nonref_genotype_count(v) for v in overlapping_variants]
if calculator.all_variants_compatible(nonref_counts):
logging.info('Overlapping variants are naturally compatible: %s',
overlapping_variants)
for variant in overlapping_variants:
yield variant
return
# The actual genotype calls produce an inconsistent haplotype. If the number
# of affected variants is "too large", avoid processing since this is an
# exponential process.
if len(overlapping_variants) > _MAX_OVERLAPPING_VARIANTS_TO_RESOLVE:
logging.warning(
'Overlapping variants are not naturally compatible, and there are too '
'many to exhaustively search (%s). Returning variants without '
'modification, beginning with %s.', len(overlapping_variants),
overlapping_variants[0])
for variant in overlapping_variants:
yield variant
return
# Otherwise, the actual genotype calls are incompatible. Since the genotype
# likelihoods are generally well-calibrated, we examine all configurations of
# genotypes that create compatible haplotypes and retain the single
# configuration with the highest joint likelihood across all variants as the
# proposed genotype assignment. Separately, we rescale the likelihood of each
# individual variant using only the valid genotype configurations. If the
# results are concordant (i.e., the genotype predicted by the marginal
# likelihood for each variant is the same as the genotype predicted when
# maximizing the joint likelihood across all variants), we return variants
# with those calls and the rescaled likelihoods. Otherwise, we log a warning
# and emit the original (incompatible) variants.
#
# For example, a biallelic deletion with probabilities of homref, het, homalt
# = 0.01, 0.9, 0.09 and inside it a biallelic SNP with probs 0.02, 0.48, 0.5.
# Naively this would be called as a heterozygous indel and a homozygous SNP,
# which is impossible as there are three total alternate genotypes. The
# algorithm does the following:
#
# Indel SNP Joint prob
# 0/0 0/0 0.01 * 0.02 = 0.0002
# 0/0 0/1 0.01 * 0.48 = 0.0048
# 0/0 1/1 0.01 * 0.50 = 0.0050
# 0/1 0/0 0.90 * 0.02 = 0.0180
# 0/1 0/1 0.90 * 0.48 = 0.4320*
# 0/1 1/1 <invalid> = 0
# 1/1 0/0 0.09 * 0.02 = 0.0018
# 1/1 0/1 <invalid> = 0
# 1/1 1/1 <invalid> = 0
#
# So using the highest joint likelihood, we predict het indel and het SNP.
#
# The marginal probability of each genotype for the indel is:
# 0/0: 0.0002 + 0.0048 + 0.0050 = 0.01
# 0/1: 0.0180 + 0.4320 = 0.45
# 1/1: 0.0018 = 0.0018
#
# which after normalizing to sum to 1 is roughly 0.022, 0.974, 0.004.
# The marginal probability for the SNP, after performing similar
# calculations, is 0.043, 0.946, 0.011. So the marginals also predict a het
# indel and a het SNP. Since the two calculations agree, we use this
# genotype call and modified likelihoods.
#
# First, we find all non-reference count configurations that are compatible.
# This represents each variant solely based on its number of non-reference
# genotypes, and assumes that variants are compatible if the total number of
# non-reference genotypes at a single position is at most two. By using
# non-reference counts, we avoid testing multiple allele configurations that
# will return the same result (e.g. a variant with two possible alternate
# alleles has three allele configurations that are homozygous alternate
# [1/1, 1/2, 2/2] and either all or none of them will be valid depending on
# the variants it interacts with).
valid_nonref_count_configurations = [
conf for conf in itertools.product(
[0, 1, 2], repeat=len(overlapping_variants))
if calculator.all_variants_compatible(conf)
]
# Next, we find the single compatible variant assignment with the individually
# highest likelihood and track the total likelihood distributed to all variant
# genotypes.
likelihood_aggregators = [
_LikelihoodAggregator(len(v.alternate_bases))
for v in overlapping_variants
]
most_likely_allele_indices_config = None
most_likely_likelihood = None
for nonref_count_config in valid_nonref_count_configurations:
for allele_indices_config in _get_all_allele_indices_configurations(
overlapping_variants, nonref_count_config):
config_likelihood = _allele_indices_configuration_likelihood(
overlapping_variants, allele_indices_config)
if (most_likely_likelihood is None or
config_likelihood > most_likely_likelihood):
most_likely_likelihood = config_likelihood
most_likely_allele_indices_config = allele_indices_config
for aggregator, allele_indices in zip(likelihood_aggregators,
allele_indices_config):
aggregator.add(allele_indices, config_likelihood)
marginal_allele_indices_config = tuple(
agg.most_likely_allele_indices() for agg in likelihood_aggregators)
if marginal_allele_indices_config == most_likely_allele_indices_config:
logging.info(
'Overlapping variants are not naturally compatible, but the genotype '
'configuration with the most likely joint likelihood is the same as '
'that from the scaled marginal likelihoods: %s',
overlapping_variants[0])
# Collapse the probabilities of all configurations to a single GL for each
# allele, independently for each variant.
scaled_gls = [agg.scaled_likelihoods() for agg in likelihood_aggregators]
for variant, allele_indices, gls in zip(
overlapping_variants, most_likely_allele_indices_config, scaled_gls):
newvariant = copy.deepcopy(variant)
call = variant_utils.only_call(newvariant)
call.genotype[:] = allele_indices
call.genotype_likelihood[:] = gls
yield newvariant
else:
logging.warning(
'Overlapping variants are not naturally compatible, and the genotype '
'configuration with the most likely joint likelihood is different from '
'that using the scaled marginal likelihoods: %s',
overlapping_variants[0])
# redacted
for variant in overlapping_variants:
yield variant
def read_squad_examples(input_file, is_training, version_2_with_negative):
"""Read a SQuAD json file into a list of SquadExample."""
with tf.io.gfile.GFile(input_file, "r") as reader:
input_data = json.load(reader)["data"]
def is_whitespace(c):
if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
return True
return False
examples = []
for entry in input_data:
for paragraph in entry["paragraphs"]:
paragraph_text = paragraph["context"]
doc_tokens = []
char_to_word_offset = []
prev_is_whitespace = True
for c in paragraph_text:
if is_whitespace(c):
prev_is_whitespace = True
else:
if prev_is_whitespace:
doc_tokens.append(c)
else:
doc_tokens[-1] += c
prev_is_whitespace = False
char_to_word_offset.append(len(doc_tokens) - 1)
for qa in paragraph["qas"]:
qas_id = qa["id"]
question_text = qa["question"]
start_position = None
end_position = None
orig_answer_text = None
is_impossible = False
if is_training:
if version_2_with_negative:
is_impossible = qa["is_impossible"]
if (len(qa["answers"]) != 1) and (not is_impossible):
raise ValueError(
"For training, each question should have exactly 1 answer.")
if not is_impossible:
answer = qa["answers"][0]
orig_answer_text = answer["text"]
answer_offset = answer["answer_start"]
answer_length = len(orig_answer_text)
start_position = char_to_word_offset[answer_offset]
end_position = char_to_word_offset[answer_offset + answer_length -
1]
# Only add answers where the text can be exactly recovered from the
# document. If this CAN'T happen it's likely due to weird Unicode
# stuff so we will just skip the example.
#
# Note that this means for training mode, every example is NOT
# guaranteed to be preserved.
actual_text = " ".join(
doc_tokens[start_position:(end_position + 1)])
cleaned_answer_text = " ".join(
tokenization.whitespace_tokenize(orig_answer_text))
if actual_text.find(cleaned_answer_text) == -1:
logging.warning("Could not find answer: '%s' vs. '%s'",
actual_text, cleaned_answer_text)
continue
else:
start_position = -1
end_position = -1
orig_answer_text = ""
example = SquadExample(
qas_id=qas_id,
question_text=question_text,
doc_tokens=doc_tokens,
orig_answer_text=orig_answer_text,
start_position=start_position,
end_position=end_position,
is_impossible=is_impossible)
examples.append(example)
return examples