def top_p_logits(logits, p):
"""Nucleus sampling"""
batch, num = logits.shape.as_list()
sorted_logits = tf.sort(logits, direction="DESCENDING", axis=-1)
cumulative_probs = tf.cumsum(tf.nn.softmax(sorted_logits, axis=-1),
axis=-1)
indices = tf.stack(
[
tf.range(0, batch),
# number of indices to include
tf.minimum(
tf.reduce_sum(tf.cast(cumulative_probs < p, tf.int32),
axis=-1), num),
],
axis=-1,
)
min_values = tf.gather_nd(sorted_logits, indices)
return tf.where(
logits < min_values,
tf.ones_like(logits) * -1e10,
logits,
)
def extract_random_video_patch(videos, num_frames=-1):
"""For every video, extract a random consecutive patch of num_frames.
Args:
videos: 5-D Tensor, (NTHWC)
num_frames: Integer, if -1 then the entire video is returned.
Returns:
video_patch: 5-D Tensor, (NTHWC) with T = num_frames.
Raises:
ValueError: If num_frames is greater than the number of total frames in
the video.
"""
if num_frames == -1:
return videos
batch_size, num_total_frames, h, w, c = common_layers.shape_list(videos)
if num_total_frames < num_frames:
raise ValueError("Expected num_frames <= %d, got %d" %
(num_total_frames, num_frames))
# Randomly choose start_inds for each video.
frame_start = tf.random_uniform(shape=(batch_size, ),
minval=0,
maxval=num_total_frames - num_frames + 1,
dtype=tf.int32)
# [start[0], start[0] + 1, ... start[0] + num_frames - 1] + ...
# [start[batch_size-1], ... start[batch_size-1] + num_frames - 1]
range_inds = tf.expand_dims(tf.range(num_frames), axis=0)
frame_inds = range_inds + tf.expand_dims(frame_start, axis=1)
frame_inds = tf.reshape(frame_inds, [-1])
# [0]*num_frames + [1]*num_frames + ... [batch_size-1]*num_frames
batch_inds = tf.expand_dims(tf.range(batch_size), axis=1)
batch_inds = tf.tile(batch_inds, [1, num_frames])
batch_inds = tf.reshape(batch_inds, [-1])
gather_inds = tf.stack((batch_inds, frame_inds), axis=1)
video_patches = tf.gather_nd(videos, gather_inds)
return tf.reshape(video_patches, (batch_size, num_frames, h, w, c))
def sparse_batched_matmul(x, y):
"""Batch multiply sparse tensor x with dense tensor y.
Args:
x: <tf.float32>[B, M, N] SparseTensor.
y: <tf.float32>[B, N, K] DenseTensor.
Returns:
<tf.float32>[B, M, K] DenseTensor.
"""
sp_indices = x.indices
sp_values = x.values
batch_size, num_row = x.dense_shape[0], x.dense_shape[1]
num_col = tf.cast(tf.shape(y)[2], tf.int64)
# Number of non-zero entries.
num_nz = tf.cast(tf.shape(sp_indices)[0], tf.int64)
# Fetch relevant values from y.
# <tf.float32>[num_nz, num_col]
lookup = tf.gather_nd(y, tf.stack(
[sp_indices[:, 0], sp_indices[:, 2]], axis=1))
# Reshape first two dimensions of x into a new SparseTensor of rank 2.
# <tf.int32>[num_nz, 2]
x_i = tf.stack([sp_indices[:, 0] * num_row + sp_indices[:, 1],
tf.cast(tf.range(num_nz), tf.int64)], axis=1)
# <tf.float32>[batch_size * num_row, num_nz]
sparse_2d = tf.SparseTensor(x_i, sp_values,
tf.stack([batch_size * num_row, num_nz], 0))
# Multiply the new sparse tensor with the gathered values.
# <tf.float32>[batch_size * num_row, num_col]
dense_2d = tf.sparse_tensor_dense_matmul(sparse_2d, lookup)
# Reshape back [batch_size, num_row, num_col]
dense_3d = tf.reshape(dense_2d, tf.stack([batch_size, num_row, num_col], 0))
return dense_3d
def body(s, alignments, idx):
"""Runs over s=[max_path-1,..., 1] (along the alignment)
:param int s: path index
:param tf.TensorArray alignments: [T+U] * (V,)
:param tf.Tensor idx: (B, 2) -> [from_u, symbol/blank]
"""
backtrack = bt_ta.read(s) # (B, U, 2)
init_u = tf.where(
tf.greater_equal(s, input_lengths + label_lengths - 1),
label_lengths, # if we are at the end of some path (or behind)
idx[:, 0] # if we are within a path, continue backtracking.
)
backtrack_indices = tf.stack([tf.range(n_batch), init_u],
axis=-1) # (B, 2)
idx = tf.gather_nd(backtrack, backtrack_indices)
align_write = tf.where(
tf.less_equal(s, input_lengths + label_lengths),
idx[:, 1], # within alignment
tf.ones((n_batch, ), dtype=tf.int32) *
blank_index) # outside, assume blank
alignments = alignments.write(s, align_write)
return s - 1, alignments, idx
def _q_tp1_doubleq(q_func,
state_p1,
batch_size,
action_size,
num_samples,
random_actions):
"""Q(s_p1, a_p1) via Double Q Learning with stochastic sampling.
Args:
q_func: Python function that takes in state, action, scope as input
and returns Q(state, action) and intermediate endpoints dictionary.
state_p1: Tensor (potentially any dtype) representing next .
batch_size: tf.Tensor (dtype=tf.int32) or int representing the minibatch
size of the state tensors.
action_size: (int) Size of continuous actio space.
num_samples: (int) Number of action samples for each minibatch element.
random_actions: tf.Tensor (dtype=tf.float32) of candidate actions.
Returns:
Tuple (q_tp1_best, end_points_tp1). See _get_q_tp1 docs for description.
"""
# Target Q values at t+1, for action selected by online net.
q_tp1_using_online_net, end_points_tp1 = q_func(
state_p1, random_actions, scope='q_func', reuse=True)
q_tp1_using_online_net_2d = tf.reshape(
q_tp1_using_online_net, (batch_size, num_samples))
q_tp1_indices_using_online_net = tf.argmax(q_tp1_using_online_net_2d, 1)
random_actions = tf.reshape(
random_actions, (batch_size, num_samples, action_size))
batch_indices = tf.cast(tf.range(batch_size), tf.int64)
indices = tf.stack([batch_indices, q_tp1_indices_using_online_net], axis=1)
# For each batch item, slice into the num_samples,
# action subarray using the corresponding to yield the chosen action.
q_tp1_best_action = tf.gather_nd(random_actions, indices)
q_tp1_best, end_points_tp1 = q_func(
state_p1, q_tp1_best_action, scope='target_q_func')
return q_tp1_best, end_points_tp1
def _build_target_distribution(self):
batch_size = tf.shape(self._replay.rewards)[0]
# size of rewards: batch_size x 1
rewards = self._replay.rewards[:, None]
# size of tiled_support: batch_size x num_atoms
is_terminal_multiplier = 1. - tf.cast(self._replay.terminals,
tf.float32)
# Incorporate terminal state to discount factor.
# size of gamma_with_terminal: batch_size x 1
gamma_with_terminal = self.cumulative_gamma * is_terminal_multiplier
gamma_with_terminal = gamma_with_terminal[:, None]
# size of next_qt_argmax: 1 x batch_size
next_qt_argmax = tf.argmax(
self._replay_next_target_net_outputs.q_values, axis=1)[:, None]
batch_indices = tf.range(tf.to_int64(batch_size))[:, None]
# size of next_qt_argmax: batch_size x 2
batch_indexed_next_qt_argmax = tf.concat(
[batch_indices, next_qt_argmax], axis=1)
# size of next_logits (next quantiles): batch_size x num_atoms
next_logits = tf.gather_nd(self._replay_next_target_net_outputs.logits,
batch_indexed_next_qt_argmax)
return rewards + gamma_with_terminal * next_logits
def _getc6(ZiZj, nci, ncj, c6ab=d3_c6ab, k3=d3_k3):
'''
interpolate c6
'''
#gather the relevant entries from the table
c6ab_ = tf.cast(tf.gather_nd(c6ab, ZiZj),nci.dtype)
#calculate c6 coefficients
c6mem = -1.0e99*tf.ones_like(nci)
r_save = 1.0e99*tf.ones_like(nci)
rsum = tf.zeros_like(nci)
csum = tf.zeros_like(nci)
for i in range(d3_maxc):
for j in range(d3_maxc):
cn0 = c6ab_[:,i,j,0]
cn1 = c6ab_[:,i,j,1]
cn2 = c6ab_[:,i,j,2]
r = (cn1-nci)**2 + (cn2-ncj)**2
r_save = tf.compat.v1.where(r < r_save, r, r_save)
c6mem = tf.compat.v1.where(r < r_save, cn0, c6mem)
tmp1 = tf.exp(k3*r)
rsum += tf.compat.v1.where(cn0 > 0.0, tmp1, tf.zeros_like(tmp1))
csum += tf.compat.v1.where(cn0 > 0.0, tmp1*cn0, tf.zeros_like(tmp1))
c6 = tf.compat.v1.where(rsum > 0.0, csum/rsum, c6mem)
return c6
def warp_img_to_uv(img_attrs, ver_proj_xy, tri_v, tri_vt, vt_list, uv_size):
""" unwrap input images to UV space """
c = img_attrs.get_shape().as_list()[-1]
uv_position = warp_ver_to_uv(ver_proj_xy, tri_v, tri_vt, vt_list, uv_size)
# build batch uv position
batch_size = img_attrs.get_shape().as_list()[0]
batch_ind = tf.reshape(
tf.range(batch_size),
[
batch_size,
1,
1,
1,
],
)
batch_ind = tf.tile(batch_ind, [1, uv_size, uv_size, 1])
uv_position = tf.cast(uv_position, tf.int32)
batch_uv_pos = tf.concat([batch_ind, uv_position], axis=-1)
batch_uv_pos = tf.reshape(batch_uv_pos, [batch_size, -1, 3])
uv_map = tf.gather_nd(img_attrs, batch_uv_pos)
uv_map = tf.reshape(uv_map, [batch_size, uv_size, uv_size, c])
return uv_map
def add_metric_fn_inputs(params, cls_outputs, box_outputs, metric_fn_inputs):
"""Selects top-k predictions and adds the selected to metric_fn_inputs.
Args:
params: a parameter dictionary that includes `min_level`, `max_level`,
`batch_size`, and `num_classes`.
cls_outputs: an OrderDict with keys representing levels and values
representing logits in [batch_size, height, width, num_anchors].
box_outputs: an OrderDict with keys representing levels and values
representing box regression targets in
[batch_size, height, width, num_anchors * 4].
metric_fn_inputs: a dictionary that will hold the top-k selections.
"""
cls_outputs_all = []
box_outputs_all = []
# Concatenates class and box of all levels into one tensor.
for level in range(params['min_level'], params['max_level'] + 1):
cls_outputs_all.append(
tf.reshape(cls_outputs[level],
[params['batch_size'], -1, params['num_classes']]))
box_outputs_all.append(
tf.reshape(box_outputs[level], [params['batch_size'], -1, 4]))
cls_outputs_all = tf.concat(cls_outputs_all, 1)
box_outputs_all = tf.concat(box_outputs_all, 1)
# cls_outputs_all has a shape of [batch_size, N, num_classes] and
# box_outputs_all has a shape of [batch_size, N, 4]. The batch_size here
# is per-shard batch size. Recently, top-k on TPU supports batch
# dimension (b/67110441), but the following function performs top-k on
# each sample.
cls_outputs_all_after_topk = []
box_outputs_all_after_topk = []
indices_all = []
classes_all = []
for index in range(params['batch_size']):
cls_outputs_per_sample = cls_outputs_all[index]
box_outputs_per_sample = box_outputs_all[index]
cls_outputs_per_sample_reshape = tf.reshape(cls_outputs_per_sample,
[-1])
_, cls_topk_indices = tf.nn.top_k(cls_outputs_per_sample_reshape,
k=anchors.MAX_DETECTION_POINTS)
# Gets top-k class and box scores.
indices = tf.div(cls_topk_indices, params['num_classes'])
classes = tf.mod(cls_topk_indices, params['num_classes'])
cls_indices = tf.stack([indices, classes], axis=1)
cls_outputs_after_topk = tf.gather_nd(cls_outputs_per_sample,
cls_indices)
cls_outputs_all_after_topk.append(cls_outputs_after_topk)
box_outputs_after_topk = tf.gather_nd(box_outputs_per_sample,
tf.expand_dims(indices, 1))
box_outputs_all_after_topk.append(box_outputs_after_topk)
indices_all.append(indices)
classes_all.append(classes)
# Concatenates via the batch dimension.
cls_outputs_all_after_topk = tf.stack(cls_outputs_all_after_topk, axis=0)
box_outputs_all_after_topk = tf.stack(box_outputs_all_after_topk, axis=0)
indices_all = tf.stack(indices_all, axis=0)
classes_all = tf.stack(classes_all, axis=0)
metric_fn_inputs['cls_outputs_all'] = cls_outputs_all_after_topk
metric_fn_inputs['box_outputs_all'] = box_outputs_all_after_topk
metric_fn_inputs['indices_all'] = indices_all
metric_fn_inputs['classes_all'] = classes_all
def assign_and_sample_proposals(proposed_boxes,
gt_boxes,
gt_classes,
gt_attributes,
num_samples_per_image=512,
mix_gt_boxes=True,
fg_fraction=0.25,
fg_iou_thresh=0.5,
bg_iou_thresh_hi=0.5,
bg_iou_thresh_lo=0.0):
"""Assigns the proposals with groundtruth classes and performs subsmpling.
Given `proposed_boxes`, `gt_boxes`, `gt_classes` and `gt_attributes`, the
function uses the following algorithm to generate the final
`num_samples_per_image` RoIs.
1. Calculates the IoU between each proposal box and each gt_boxes.
2. Assigns each proposed box with a groundtruth class and box by choosing
the largest IoU overlap.
3. Samples `num_samples_per_image` boxes from all proposed boxes, and
returns box_targets, class_targets, and RoIs.
Args:
proposed_boxes: a tensor of shape of [batch_size, N, 4]. N is the number
of proposals before groundtruth assignment. The last dimension is the
box coordinates w.r.t. the scaled images in [ymin, xmin, ymax, xmax]
format.
gt_boxes: a tensor of shape of [batch_size, MAX_NUM_INSTANCES, 4].
The coordinates of gt_boxes are in the pixel coordinates of the scaled
image. This tensor might have padding of values -1 indicating the invalid
box coordinates.
gt_classes: a tensor with a shape of [batch_size, MAX_NUM_INSTANCES]. This
tensor might have paddings with values of -1 indicating the invalid
classes.
gt_attributes: a tensor with a shape of [batch_size, MAX_NUM_INSTANCES,
num_attributes]. This tensor might have paddings with values of -1
indicating the invalid attributes.
num_samples_per_image: an integer represents RoI minibatch size per image.
mix_gt_boxes: a bool indicating whether to mix the groundtruth boxes before
sampling proposals.
fg_fraction: a float represents the target fraction of RoI minibatch that
is labeled foreground (i.e., class > 0).
fg_iou_thresh: a float represents the IoU overlap threshold for an RoI to be
considered foreground (if >= fg_iou_thresh).
bg_iou_thresh_hi: a float represents the IoU overlap threshold for an RoI to
be considered background (class = 0 if overlap in [LO, HI)).
bg_iou_thresh_lo: a float represents the IoU overlap threshold for an RoI to
be considered background (class = 0 if overlap in [LO, HI)).
Returns:
sampled_rois: a tensor of shape of [batch_size, K, 4], representing the
coordinates of the sampled RoIs, where K is the number of the sampled
RoIs, i.e. K = num_samples_per_image.
sampled_gt_boxes: a tensor of shape of [batch_size, K, 4], storing the
box coordinates of the matched groundtruth boxes of the samples RoIs.
sampled_gt_classes: a tensor of shape of [batch_size, K], storing the
classes of the matched groundtruth boxes of the sampled RoIs.
sampled_gt_attributes: a tensor of shape of [batch_size, K,
num_attributes], storing the attributes of the matched groundtruth
attributes of the sampled RoIs.
sampled_gt_indices: a tensor of shape of [batch_size, K], storing the
indices of the sampled groudntruth boxes in the original `gt_boxes`
tensor, i.e. gt_boxes[sampled_gt_indices[:, i]] = sampled_gt_boxes[:, i].
"""
with tf.name_scope('sample_proposals'):
if mix_gt_boxes:
boxes = tf.concat([proposed_boxes, gt_boxes], axis=1)
else:
boxes = proposed_boxes
(matched_gt_boxes, matched_gt_classes, matched_gt_attributes,
matched_gt_indices, matched_iou,
_) = box_matching(boxes, gt_boxes, gt_classes, gt_attributes)
positive_match = tf.greater(matched_iou, fg_iou_thresh)
negative_match = tf.logical_and(
tf.greater_equal(matched_iou, bg_iou_thresh_lo),
tf.less(matched_iou, bg_iou_thresh_hi))
ignored_match = tf.less(matched_iou, 0.0)
# re-assign negatively matched boxes to the background class.
matched_gt_classes = tf.where(negative_match,
tf.zeros_like(matched_gt_classes),
matched_gt_classes)
matched_gt_indices = tf.where(negative_match,
tf.zeros_like(matched_gt_indices),
matched_gt_indices)
sample_candidates = tf.logical_and(
tf.logical_or(positive_match, negative_match),
tf.logical_not(ignored_match))
sampler = (
balanced_positive_negative_sampler.BalancedPositiveNegativeSampler(
positive_fraction=fg_fraction, is_static=True))
batch_size, _ = sample_candidates.get_shape().as_list()
sampled_indicators = []
for i in range(batch_size):
sampled_indicator = sampler.subsample(sample_candidates[i],
num_samples_per_image,
positive_match[i])
sampled_indicators.append(sampled_indicator)
sampled_indicators = tf.stack(sampled_indicators)
_, sampled_indices = tf.nn.top_k(tf.cast(sampled_indicators,
dtype=tf.int32),
k=num_samples_per_image,
sorted=True)
sampled_indices_shape = tf.shape(sampled_indices)
batch_indices = (
tf.expand_dims(tf.range(sampled_indices_shape[0]), axis=-1) *
tf.ones([1, sampled_indices_shape[-1]], dtype=tf.int32))
gather_nd_indices = tf.stack([batch_indices, sampled_indices], axis=-1)
sampled_rois = tf.gather_nd(boxes, gather_nd_indices)
sampled_gt_boxes = tf.gather_nd(matched_gt_boxes, gather_nd_indices)
sampled_gt_classes = tf.gather_nd(matched_gt_classes,
gather_nd_indices)
sampled_gt_attributes = tf.gather_nd(matched_gt_attributes,
gather_nd_indices)
sampled_gt_indices = tf.gather_nd(matched_gt_indices,
gather_nd_indices)
return (sampled_rois, sampled_gt_boxes, sampled_gt_classes,
sampled_gt_attributes, sampled_gt_indices)
def build_bert_inputs(example):
"""Convert example <Tensor [30, 70]> into bert inputs."""
k_size = FLAGS.k_size
CLS_ID = tf.constant([101], dtype=tf.int64) # pylint: disable=invalid-name
SEP_ID = tf.constant([102], dtype=tf.int64) # pylint: disable=invalid-name
max_len = tf.constant([FLAGS.max_para_length])
context_size = tf.constant([FLAGS.context_size])
intermediate_examples_tensor = tf.reduce_sum(tf.abs(example), 1)
examples_zero_vector = tf.zeros(shape=(1, 1), dtype=tf.int64)
examples_bool_mask = tf.squeeze(
tf.not_equal(intermediate_examples_tensor, examples_zero_vector))
paragraph_len = tf.reduce_sum(tf.cast(examples_bool_mask, tf.int32))
start = tf.random.uniform([1],
0,
tf.reshape(paragraph_len, []) -
tf.reshape(context_size, []) + 1,
dtype=tf.int32)
# Slice the document into the before, after and context.
# Discard the zero padding.
sizes = tf.squeeze(
tf.concat([[
start, context_size, paragraph_len - context_size - start,
max_len - paragraph_len
]], 0))
before, context, after, _ = tf.split(example, sizes, axis=0)
# Gather the context removing zero padding at end of sentences.
non_zeros = tf.where(tf.not_equal(context, tf.zeros_like(context)))
context_gathered = tf.gather_nd(context, non_zeros)
# Flip before so we select the 4 sentences closest to target
before = tf.reverse(before, axis=[0])
# pad both to longer than needed
paddings = tf.constant([[0, 8], [0, 0]])
before = tf.pad(before, paddings)
after = tf.pad(after, paddings)
# Extend targets to 3 sentences
# pad both
before_minus_one = before[1:][:k_size]
before_minus_two = before[2:][:k_size]
after_plus_one = after[1:][:k_size]
after_plus_two = after[2:][:k_size]
before = before[:k_size]
after = after[:k_size]
before = tf.concat([before_minus_two, before_minus_one, before], axis=1)
after = tf.concat([after, after_plus_one, after_plus_two], axis=1)
############################################################################
# These 8 sentences are the 8 surrounding targets. Some are padding.
targets = tf.concat([before, after], axis=0)
# Remove the padding from the sourrounding sentences
# Eg. if context starts at beginning of paragraph, before is all padding
intermediate_tensor = tf.reduce_sum(tf.abs(targets), 1)
zero_vector = tf.zeros(shape=(1, 1), dtype=tf.int64)
bool_mask = tf.squeeze(tf.not_equal(intermediate_tensor, zero_vector))
bool_mask.set_shape([None])
targets = tf.boolean_mask(targets, bool_mask)
# Randomly select 4 targets
# We will also select the label_types for each selected target
indices = tf.range(0, limit=tf.shape(targets)[0], dtype=tf.int32)
shuffled_indices = tf.random.shuffle(indices)[:k_size]
targets = tf.gather(targets, shuffled_indices)
if k_size == 4:
full_labels = tf.concat([tf.range(3, -1, -1), tf.range(4, 8)], axis=0)
elif k_size == 3:
full_labels = tf.concat([tf.range(2, -1, -1), tf.range(3, 6)], axis=0)
elif k_size == 2:
full_labels = tf.concat([tf.range(1, -1, -1), tf.range(2, 4)], axis=0)
elif k_size == 1:
full_labels = tf.concat([tf.range(0, -1, -1), tf.range(1, 2)], axis=0)
label_types = tf.boolean_mask(full_labels, bool_mask)
label_types = tf.gather(label_types, shuffled_indices)
# create inputs
bert_inputs = []
input_masks = []
segment_ids = []
# make context
ctx_segment_id = tf.concat([
tf.zeros_like(CLS_ID, dtype=tf.int64),
tf.zeros_like(context_gathered),
tf.zeros_like(SEP_ID, dtype=tf.int64)
],
axis=0)
ctx_segment_id = pad_and_cut(ctx_segment_id, FLAGS.max_seq_length)
segment_ids.append(ctx_segment_id)
new_ctx_input = tf.concat([CLS_ID, context_gathered, SEP_ID], axis=0)
ctx_input_mask = tf.ones_like(new_ctx_input)
ctx_input_mask = pad_and_cut(ctx_input_mask, FLAGS.max_seq_length)
input_masks.append(ctx_input_mask)
padded_new_ctx_input = pad_and_cut(new_ctx_input, FLAGS.max_seq_length)
bert_inputs.append(padded_new_ctx_input)
for i in range(k_size):
target_non_zero = tf.where(
tf.not_equal(targets[i], tf.zeros_like(targets[i])))
targets_stripped = tf.gather_nd(targets[i], target_non_zero)
if FLAGS.include_context:
segment_id = tf.concat([
tf.zeros_like(CLS_ID, dtype=tf.int64),
tf.zeros_like(context_gathered),
tf.zeros_like(SEP_ID, dtype=tf.int64),
tf.ones_like(targets_stripped),
tf.ones_like(SEP_ID, dtype=tf.int64)
],
axis=0)
else:
segment_id = tf.concat([
tf.zeros_like(CLS_ID, dtype=tf.int64),
tf.zeros_like(targets_stripped),
tf.zeros_like(SEP_ID, dtype=tf.int64)
],
axis=0)
segment_id = pad_and_cut(segment_id, FLAGS.max_seq_length)
segment_ids.append(segment_id)
if FLAGS.include_context:
new_input = tf.concat(
[CLS_ID, context_gathered, SEP_ID, targets_stripped, SEP_ID], axis=0)
else:
new_input = tf.concat([CLS_ID, targets_stripped, SEP_ID], axis=0)
input_mask = tf.ones_like(new_input)
input_mask = pad_and_cut(input_mask, FLAGS.max_seq_length)
input_masks.append(input_mask)
padded_new_input = pad_and_cut(new_input, FLAGS.max_seq_length)
bert_inputs.append(padded_new_input)
bert_inputs = tf.stack(bert_inputs, axis=0)
input_masks = tf.stack(input_masks, axis=0)
segment_ids = tf.stack(segment_ids, axis=0)
out = Outputs_And_Context(bert_inputs, input_masks, segment_ids, label_types,
context_gathered)
return out
b_3_p = tf.Variable(tf.truncated_normal([n_hidden_3], stddev=0.1))
b_4_p = tf.Variable(tf.truncated_normal([n_output], stddev=0.1))
layer_1_p = tf.nn.relu(tf.add(tf.matmul(x_p, w_1_p), b_1_p))
layer_1_p_b = tf.layers.batch_normalization(layer_1_p)
layer_2_p = tf.nn.relu(tf.add(tf.matmul(layer_1_p_b, w_2_p), b_2_p))
layer_2_p_b = tf.layers.batch_normalization(layer_2_p)
layer_3_p = tf.nn.relu(tf.add(tf.matmul(layer_2_p_b, w_3_p), b_3_p))
layer_3_p_b = tf.layers.batch_normalization(layer_3_p)
y_p = tf.nn.relu(tf.add(tf.matmul(layer_3_p_b, w_4_p), b_4_p))
g_target_q_idx = tf.placeholder('int32', [None, None], 'output_idx')
target_q_with_idx = tf.gather_nd(y_p, g_target_q_idx)
init = tf.global_variables_initializer()
saver = tf.train.Saver()
def predict(sess, s_t, ep, test_ep=False):
n_power_levels = len(env.V2V_power_dB_List)
if np.random.rand() < ep and not test_ep:
pred_action = np.random.randint(n_RB * n_power_levels)
else:
pred_action = sess.run(g_q_action, feed_dict={x: [s_t]})[0]
return pred_action
def add_metric_fn_inputs(params,
cls_outputs,
box_outputs,
metric_fn_inputs,
max_detection_points=anchors.MAX_DETECTION_POINTS):
"""Selects top-k predictions and adds the selected to metric_fn_inputs.
Args:
params: a parameter dictionary that includes `min_level`, `max_level`,
`batch_size`, and `num_classes`.
cls_outputs: an OrderDict with keys representing levels and values
representing logits in [batch_size, height, width, num_anchors].
box_outputs: an OrderDict with keys representing levels and values
representing box regression targets in [batch_size, height, width,
num_anchors * 4].
metric_fn_inputs: a dictionary that will hold the top-k selections.
max_detection_points: an integer specifing the maximum detection points to
keep before NMS. Keep all anchors if max_detection_points <= 0.
"""
batch_size = params['batch_size']
num_classes = params['num_classes']
cls_outputs_all = []
box_outputs_all = []
# Concatenates class and box of all levels into one tensor.
for level in range(params['min_level'], params['max_level'] + 1):
if params['data_format'] == 'channels_first':
cls_outputs[level] = tf.transpose(cls_outputs[level], [0, 2, 3, 1])
box_outputs[level] = tf.transpose(box_outputs[level], [0, 2, 3, 1])
cls_outputs_all.append(
tf.reshape(cls_outputs[level], [batch_size, -1, num_classes]))
box_outputs_all.append(
tf.reshape(box_outputs[level], [batch_size, -1, 4]))
cls_outputs_all = tf.concat(cls_outputs_all, 1)
box_outputs_all = tf.concat(box_outputs_all, 1)
if max_detection_points > 0:
# Prune anchors and detections to only keep max_detection_points.
# Due to some issues, top_k is currently slow in graph model.
cls_outputs_all_reshape = tf.reshape(cls_outputs_all, [batch_size, -1])
_, cls_topk_indices = tf.math.top_k(cls_outputs_all_reshape,
k=max_detection_points,
sorted=False)
indices = cls_topk_indices // num_classes
classes = cls_topk_indices % num_classes
cls_indices = tf.stack([indices, classes], axis=2)
cls_outputs_all_after_topk = tf.gather_nd(cls_outputs_all,
cls_indices,
batch_dims=1)
box_outputs_all_after_topk = tf.gather_nd(box_outputs_all,
tf.expand_dims(indices, 2),
batch_dims=1)
else:
# Keep all anchors, but for each anchor, just keep the max probablity for
# each class.
cls_outputs_idx = tf.math.argmax(cls_outputs_all,
axis=-1,
output_type=tf.int32)
num_anchors = cls_outputs_all.shape[1]
classes = cls_outputs_idx
indices = tf.tile(tf.expand_dims(tf.range(num_anchors), axis=0),
[batch_size, 1])
cls_outputs_all_after_topk = tf.reduce_max(cls_outputs_all, -1)
box_outputs_all_after_topk = box_outputs_all
metric_fn_inputs['cls_outputs_all'] = cls_outputs_all_after_topk
metric_fn_inputs['box_outputs_all'] = box_outputs_all_after_topk
metric_fn_inputs['indices_all'] = indices
metric_fn_inputs['classes_all'] = classes
def mask(config: configure_pretraining.PretrainingConfig,
inputs: pretrain_data.Inputs, mask_prob, proposal_distribution=1.0,
disallow_from_mask=None, already_masked=None):
"""Implementation of dynamic masking. The optional arguments aren't needed for
BERT/ELECTRA and are from early experiments in "strategically" masking out
tokens instead of uniformly at random.
Args:
config: configure_pretraining.PretrainingConfig
inputs: pretrain_data.Inputs containing input input_ids/input_mask
mask_prob: percent of tokens to mask
proposal_distribution: for non-uniform masking can be a [B, L] tensor
of scores for masking each position.
disallow_from_mask: a boolean tensor of [B, L] of positions that should
not be masked out
already_masked: a boolean tensor of [B, N] of already masked-out tokens
for multiple rounds of masking
Returns: a pretrain_data.Inputs with masking added
"""
# Get the batch size, sequence length, and max masked-out tokens
N = config.max_predictions_per_seq
B, L = modeling.get_shape_list(inputs.input_ids)
# Find indices where masking out a token is allowed
vocab = tokenization.FullTokenizer(
config.vocab_file, do_lower_case=config.do_lower_case).vocab
candidates_mask = _get_candidates_mask(inputs, vocab, disallow_from_mask)
# Set the number of tokens to mask out per example
num_tokens = tf.cast(tf.reduce_sum(inputs.input_mask, -1), tf.float32)
num_to_predict = tf.maximum(1, tf.minimum(
N, tf.cast(tf.round(num_tokens * mask_prob), tf.int32)))
masked_lm_weights = tf.cast(tf.sequence_mask(num_to_predict, N), tf.float32)
if already_masked is not None:
masked_lm_weights *= (1 - already_masked)
# Get a probability of masking each position in the sequence
candidate_mask_float = tf.cast(candidates_mask, tf.float32)
sample_prob = (proposal_distribution * candidate_mask_float)
sample_prob /= tf.reduce_sum(sample_prob, axis=-1, keepdims=True)
# Sample the positions to mask out
sample_prob = tf.stop_gradient(sample_prob)
sample_logits = tf.log(sample_prob)
masked_lm_positions = tf.random.categorical(
sample_logits, N, dtype=tf.int32)
masked_lm_positions *= tf.cast(masked_lm_weights, tf.int32)
# Get the ids of the masked-out tokens
shift = tf.expand_dims(L * tf.range(B), -1)
flat_positions = tf.reshape(masked_lm_positions + shift, [-1, 1])
masked_lm_ids = tf.gather_nd(tf.reshape(inputs.input_ids, [-1]),
flat_positions)
masked_lm_ids = tf.reshape(masked_lm_ids, [B, -1])
masked_lm_ids *= tf.cast(masked_lm_weights, tf.int32)
# Update the input ids
replace_with_mask_positions = masked_lm_positions * tf.cast(
tf.less(tf.random.uniform([B, N]), 0.85), tf.int32)
inputs_ids, _ = scatter_update(
inputs.input_ids, tf.fill([B, N], vocab["[MASK]"]),
replace_with_mask_positions)
return pretrain_data.get_updated_inputs(
inputs,
input_ids=tf.stop_gradient(inputs_ids),
masked_lm_positions=masked_lm_positions,
masked_lm_ids=masked_lm_ids,
masked_lm_weights=masked_lm_weights
)
def __init__(self,
dataset,
embedding_size=100,
margin=1.0,
learning_rate=1e-3,
batch_size=1024,
sampling_type='bernoulli',
patience=50):
self.train_loss_history = []
self.validation_loss_history = []
self.dataset = dataset
self.entity_count = dataset.entity_count
self.relation_count = dataset.relation_count
self.embedding_size = embedding_size
self.margin = margin
self.learning_rate = learning_rate
self.batch_size = batch_size
self.patience = patience
self.sampling_type = sampling_type
self._labels = {}
self.entity_embeddings = tensorflow.get_variable(
name="entity_embeddings",
shape=[self.entity_count, self.embedding_size])
self.relation_embeddings = tensorflow.get_variable(
name="relation_embeddings",
shape=[self.relation_count, self.embedding_size])
self.weight = tensorflow.get_variable(name="weight_matrix",
shape=[
self.embedding_size,
self.embedding_size,
self.embedding_size
])
self.inp_dropout = tensorflow.keras.layers.Dropout(rate=0.3)
self.hidden_dropout1 = tensorflow.keras.layers.Dropout(rate=0.4)
self.hidden_dropout2 = tensorflow.keras.layers.Dropout(rate=0.5)
self.head_id = tensorflow.placeholder(tensorflow.int32)
self.tail_id = tensorflow.placeholder(tensorflow.int32)
self.relation_id = tensorflow.placeholder(tensorflow.int32)
self.label = tensorflow.placeholder(tensorflow.float32)
head_embedding_vector = tensorflow.nn.embedding_lookup(
self.entity_embeddings, self.head_id)
head_embedding_vector = tensorflow.nn.l2_normalize(
head_embedding_vector, axis=1)
head_embedding_vector = self.inp_dropout(head_embedding_vector)
head_embedding_vector = tensorflow.reshape(
head_embedding_vector, [-1, 1, self.embedding_size])
relation_embedding_vector = tensorflow.nn.embedding_lookup(
self.relation_embeddings, self.relation_id)
W_mat = tensorflow.matmul(
relation_embedding_vector,
tensorflow.reshape(self.weight, [self.embedding_size, -1]))
W_mat = tensorflow.reshape(
W_mat, [-1, self.embedding_size, self.embedding_size])
W_mat = self.hidden_dropout1(W_mat)
x = tensorflow.matmul(head_embedding_vector, W_mat)
x = tensorflow.reshape(x, [-1, self.embedding_size])
x = tensorflow.nn.l2_normalize(x, axis=1)
x = self.hidden_dropout2(x)
x = tensorflow.matmul(x, self.entity_embeddings, transpose_b=True)
out = tensorflow.nn.sigmoid(x)
indices = tensorflow.concat([
tensorflow.reshape(tensorflow.range(tensorflow.size(self.head_id)),
(-1, 1)),
tensorflow.reshape(self.tail_id, (-1, 1))
],
axis=1)
predict = out
self.predict = -tensorflow.reshape(tensorflow.gather_nd(out, indices),
(-1, ))
bce = tensorflow.keras.losses.BinaryCrossentropy()
self.loss = bce(self.label, predict)
self.optimizer = tensorflow.train.AdamOptimizer(
learning_rate=self.learning_rate)
self.learn = self.optimizer.minimize(self.loss)
def real_svg_top(body_output,
unused_targets,
model_hparams,
unused_vocab_size,
hard=False):
"""Applies the Mixture Density Network on top of the LSTM outputs.
Args:
body_output: outputs from LSTM with shape [batch, seqlen, 1, hidden_size]
unused_targets: what the ground truth SVG outputted should be (unused).
model_hparams: hyper-parameters, should include num_mixture,
mix_temperature, and gauss_temperature.
unused_vocab_size: unused
hard: whether to force predict mode functionality, or return all MDN
components
Returns:
The MDN output. Could be shape [batch, seqlen, 1, 10] if in predict mode
(or hard=True) or shape [batch, seqlen, 1, 4 + 6 * num_mix * 3], in train.
"""
# mixture of gaussians for 6 args plus 4 extra states for cmds
num_mix = model_hparams.num_mixture
nout = 4 + 6 * num_mix * 3
# the 'hard' option is meant to be used if 'top' is called within body
with tf.variable_scope('real_top', reuse=tf.AUTO_REUSE):
ret = tf.layers.dense(body_output, nout, name='top')
batch_size = common_layers.shape_list(ret)[0]
if hard or model_hparams.mode == tf.estimator.ModeKeys.PREDICT:
temperature = model_hparams.mix_temperature
# apply temperature, do softmax
command = tf.identity(ret[:, :, :, :4]) / temperature
command = tf.exp(command -
tf.reduce_max(command, axis=[-1], keepdims=True))
command = command / tf.reduce_sum(
command, axis=[-1], keepdims=True)
# sample from the given probs, this is the same as get_pi_idx,
# and already returns not soft prob
command = tf.distributions.Categorical(probs=command).sample()
# this is now [batch, seq, 1], need to make it one_hot
command = tf.one_hot(command, 4)
arguments = ret[:, :, :, 4:]
# args are [batch, seq, 1, 6*3*num_mix]. want [batch * seq * 6, 3*num_mix]
arguments = tf.reshape(arguments, [-1, 3 * num_mix])
out_logmix, out_mean, out_logstd = _get_mdn_coef(arguments)
# these are [batch*seq*6, num_mix]
# apply temp to logmix
out_logmix = tf.identity(out_logmix) / temperature
out_logmix = tf.exp(
out_logmix -
tf.reduce_max(out_logmix, axis=[-1], keepdims=True))
out_logmix = out_logmix / tf.reduce_sum(
out_logmix, axis=[-1], keepdims=True)
# get_pi_idx
out_logmix = tf.distributions.Categorical(
probs=out_logmix).sample()
# should now be [batch*seq*6, 1]
out_logmix = tf.cast(out_logmix, tf.int32)
out_logmix = tf.reshape(out_logmix, [-1])
# prepare for gather
out_logmix = tf.stack([tf.range(tf.size(out_logmix)), out_logmix],
axis=-1)
chosen_mean = tf.gather_nd(out_mean, out_logmix)
chosen_logstd = tf.gather_nd(out_logstd, out_logmix)
# sample!!
rand_gaussian = (tf.random.normal(tf.shape(chosen_mean)) *
tf.sqrt(model_hparams.gauss_temperature))
arguments = chosen_mean + tf.exp(chosen_logstd) * rand_gaussian
arguments = tf.reshape(arguments, [batch_size, -1, 1, 6])
# concat with the command we picked!
ret = tf.concat([command, arguments], axis=-1)
return ret
def generate_detections_per_image_tpu(cls_outputs,
box_outputs,
anchor_boxes,
image_info,
pre_nms_num_detections=1000,
post_nms_num_detections=100,
nms_threshold=0.3,
bbox_reg_weights=(10., 10., 5., 5.)):
"""Generate the final detections per image given the model outputs.
Args:
cls_outputs: a tensor with shape [N, num_classes], which stacks class
logit outputs on all feature levels. The N is the number of total anchors
on all levels. The num_classes is the number of classes predicted by the
model. Note that the cls_outputs should be the output of softmax().
box_outputs: a tensor with shape [N, num_classes*4], which stacks box
regression outputs on all feature levels. The N is the number of total
anchors on all levels.
anchor_boxes: a tensor with shape [N, 4], which stacks anchors on all
feature levels. The N is the number of total anchors on all levels.
image_info: a tensor of shape [5] which encodes the input image's [height,
width, scale, original_height, original_width]
pre_nms_num_detections: an integer that specifies the number of candidates
before NMS.
post_nms_num_detections: an integer that specifies the number of candidates
after NMS.
nms_threshold: a float number to specify the IOU threshold of NMS.
bbox_reg_weights: a list of 4 float scalars, which are default weights on
(dx, dy, dw, dh) for normalizing bbox regression targets.
Returns:
detections: Tuple of tensors corresponding to number of valid boxes,
box coordinates, object categories for each boxes, and box scores
-- respectively.
"""
num_boxes, num_classes = cls_outputs.get_shape().as_list()
# Remove background class scores.
cls_outputs = cls_outputs[:, 1:num_classes]
top_k_scores, top_k_indices_with_classes = tf.nn.top_k(
tf.reshape(cls_outputs, [-1]), k=pre_nms_num_detections, sorted=False)
classes = tf.mod(top_k_indices_with_classes, num_classes - 1)
top_k_indices = tf.floordiv(top_k_indices_with_classes, num_classes - 1)
anchor_boxes = tf.gather(anchor_boxes, top_k_indices)
box_outputs = tf.reshape(box_outputs,
[num_boxes, num_classes, 4])[:, 1:num_classes, :]
class_indices = classes
box_outputs = tf.gather_nd(
box_outputs, tf.stack([top_k_indices, class_indices], axis=1))
# apply bounding box regression to anchors
boxes = box_utils.decode_boxes(box_outputs, anchor_boxes, bbox_reg_weights)
boxes = box_utils.clip_boxes(boxes, image_info[0], image_info[1])
list_of_all_boxes = []
list_of_all_scores = []
list_of_all_classes = []
# Skip background class.
for class_i in range(num_classes):
# Compute bitmask for the given classes.
class_i_bitmask = tf.cast(tf.equal(classes, class_i),
top_k_scores.dtype)
# This works because score is in [0, 1].
class_i_scores = top_k_scores * class_i_bitmask
# The TPU and CPU have different behaviors for
# tf.image.non_max_suppression_padded (b/116754376).
(class_i_post_nms_indices,
class_i_nms_num_valid) = tf.image.non_max_suppression_padded(
tf.to_float(boxes),
tf.to_float(class_i_scores),
post_nms_num_detections,
iou_threshold=nms_threshold,
score_threshold=0.05,
pad_to_max_output_size=True,
name='nms_detections_' + str(class_i))
class_i_post_nms_boxes = tf.gather(boxes, class_i_post_nms_indices)
class_i_post_nms_scores = tf.gather(class_i_scores,
class_i_post_nms_indices)
mask = tf.less(tf.range(post_nms_num_detections),
[class_i_nms_num_valid])
class_i_post_nms_scores = tf.where(
mask, class_i_post_nms_scores,
tf.zeros_like(class_i_post_nms_scores))
class_i_classes = tf.fill(tf.shape(class_i_post_nms_scores),
class_i + 1)
list_of_all_boxes.append(class_i_post_nms_boxes)
list_of_all_scores.append(class_i_post_nms_scores)
list_of_all_classes.append(class_i_classes)
post_nms_boxes = tf.concat(list_of_all_boxes, axis=0)
post_nms_scores = tf.concat(list_of_all_scores, axis=0)
post_nms_classes = tf.concat(list_of_all_classes, axis=0)
# sort all results.
post_nms_scores, sorted_indices = tf.nn.top_k(tf.to_float(post_nms_scores),
k=post_nms_num_detections,
sorted=True)
post_nms_boxes = tf.gather(post_nms_boxes, sorted_indices)
post_nms_classes = tf.gather(post_nms_classes, sorted_indices)
valid_mask = tf.where(tf.greater(post_nms_scores, 0),
tf.ones_like(post_nms_scores),
tf.zeros_like(post_nms_scores))
num_valid_boxes = tf.reduce_sum(valid_mask, axis=-1)
box_classes = tf.to_float(post_nms_classes)
return num_valid_boxes, post_nms_boxes, box_classes, post_nms_scores
def _build_net(self):
self.s = tf.placeholder(tf.float32, [None, self.n_features], name='s')
self.s_ = tf.placeholder(tf.float32, [None, self.n_features],
name='s_')
self.r = tf.placeholder(tf.float32, [
None,
], name='r')
self.a = tf.placeholder(tf.int32, [
None,
], name='a')
w_initializer, b_initializer = tf.random_normal_initializer(
0.0, 0.3), tf.constant_initializer(0.1)
with tf.variable_scope('eval_net'):
e1 = tf.layers.dense(self.s,
30,
tf.nn.relu,
kernel_initializer=w_initializer,
bias_initializer=b_initializer,
name='e1')
self.q_el = tf.layers.dense(e1,
10,
kernel_initializer=w_initializer,
bias_initializer=b_initializer,
name='e2')
self.q_eval = tf.layers.dense(self.q_el,
self.n_actions,
kernel_initializer=w_initializer,
bias_initializer=b_initializer,
name='e3')
with tf.variable_scope('target_net'):
t1 = tf.layers.dense(self.s_,
30,
tf.nn.relu,
kernel_initializer=w_initializer,
bias_initializer=b_initializer,
name='t1')
self.t_el = tf.layers.dense(t1,
10,
kernel_initializer=w_initializer,
bias_initializer=b_initializer,
name='t2')
self.q_next = tf.layers.dense(self.t_el,
self.n_actions,
kernel_initializer=w_initializer,
bias_initializer=b_initializer,
name='t3')
with tf.variable_scope('q_target'):
q_target = self.r + self.gamma * tf.reduce_max(
self.q_next, axis=1, name='Qmax_s_') # shape=(None, )
self.q_target = tf.stop_gradient(q_target)
with tf.variable_scope('q_eval'):
a_indices = tf.stack(
[tf.range(tf.shape(self.a)[0], dtype=tf.int32), self.a],
axis=1)
self.q_eval_wrt_a = tf.gather_nd(
params=self.q_eval, indices=a_indices) # shape=(None, )
with tf.variable_scope('loss'):
self.loss = tf.reduce_mean(
tf.squared_difference(self.q_target,
self.q_eval_wrt_a,
name='TD_error'))
with tf.variable_scope('train'):
self._train_op = tf.train.AdamOptimizer(self.lr).minimize(
self.loss)
def get_final(sequence, sequence_length, time_major=True):
"""Get the final item in a batch of sequences."""
final_index = _get_final_index(sequence_length, time_major)
return tf.gather_nd(sequence, final_index)
def geometry_augment(images, p=1, PADDING="REFLECT"):
B, real_H, real_W, C = images.shape.as_list()
#pad_size=real_H//2
#mirror_pad_images=tf.pad(images, [[0,0],[pad_size, pad_size],[pad_size, pad_size], [0, 0]],'REFLECT') #gradient wrt the padded region is unimplemented on TPUs. Constant padding can result in (possibly non-square) cutout regions at the margins of the final output, which may be desirable.
#constant_pad_images=tf.pad(images, [[0,0],[pad_size, pad_size],[pad_size, pad_size], [0, 0]],'CONSTANT')
#images=constant_pad_images#tf.stop_gradient(mirror_pad_images-constant_pad_images)+constant_pad_images
#images=tf.compat.v1.image.resize(images, [real_H*4,real_W*4])
#images=tf.stop_gradient(mirror_pad_images-constant_pad_images)+constant_pad_images
images = upsample(images)
B, H, W, C = images.shape.as_list()
DIM = H
XDIM = DIM % 2 #fix for size 331
m = tf.reshape(tf.eye(3), [1, 3, 3])
#Transform the doubled coordinates back to the original coordinates (multiply by S^-1)
m = get_scale_mat(tf.ones([B]) / 2, tf.ones([B]) / 2) @ m
#Horizontal flip
toss = tf.cast(tf.random.uniform([B]) < p, tf.float32)
m = get_scale_mat(tf.ones([B]), tf.ones([B]) - 2 * toss) @ m
#90-degree rotation
toss = tf.reshape(
tf.random.categorical(
tf.tile(tf.math.log([[1 - p + p / 4, p / 4, p / 4, p / 4]]),
[B, 1]), 1), [B])
rad1 = tf.cast(toss, tf.float32) * np.pi / 2
m = get_rot_mat(rad1) @ m
#Integer translation
toss = tf.cast(tf.random.uniform([B]) < p, tf.float32)
x_offset = tf.cast(
tf.round(toss * tf.random.uniform([B], -.125, .125) * real_W),
tf.float32)
y_offset = tf.cast(
tf.round(toss * tf.random.uniform([B], -.125, .125) * real_H),
tf.float32)
m = get_shift_mat(y_offset, x_offset) @ m
#Isotropic scaling
toss = tf.cast(tf.random.uniform([B]) < p, tf.float32)
scale = tf.exp(toss * tf.random.normal([B], 0, 0.2 * tf.math.log(2.)))
m = get_scale_mat(scale, scale) @ m
#Pre-rotation
p_rot = 1. - tf.sqrt(1. - p)
toss = tf.cast(tf.random.uniform([B]) < p_rot, tf.float32)
rad2 = toss * tf.random.uniform([B], -np.pi, np.pi)
m = get_rot_mat(rad2) @ m
#Anisotropic scaling
toss = tf.cast(tf.random.uniform([B]) < p, tf.float32)
x_scale = tf.exp(toss * tf.random.normal([B], 0, 0.2 * tf.math.log(2.)))
y_scale = 1. / x_scale
m = get_scale_mat(y_scale, x_scale) @ m
#Post-rotation
p_rot = 1. - tf.sqrt(1. - p)
toss = tf.cast(tf.reshape(tf.random.uniform([B]) < p_rot, [B]), tf.float32)
rad3 = toss * tf.random.uniform([B], -np.pi, np.pi)
m = get_rot_mat(rad3) @ m
#Fractional translation
toss = tf.cast(tf.random.uniform([B]) < p, tf.float32)
x_offset = toss * tf.random.normal([B], 0, .125) * real_W
y_offset = toss * tf.random.normal([B], 0, .125) * real_H
m = get_shift_mat(y_offset, x_offset) @ m
# Transform to the coordinates of the upsampled images (multiply by S)
m = get_scale_mat(tf.ones([B]) * 2, tf.ones([B]) * 2) @ m
# LIST DESTINATION PIXEL INDICES
#x = tf.repeat( tf.range(DIM//2,-DIM//2,-1), DIM ) # TF1.15 or above
x = tf.reshape(tf.transpose([tf.range(DIM // 2, -DIM // 2, -1)] * DIM),
[-1])
y = tf.tile(tf.range(-DIM // 2, DIM // 2), [DIM])
z = tf.ones([DIM * DIM], dtype='int32')
idx = tf.stack([x, y, z])
# ROTATE DESTINATION PIXELS ONTO ORIGIN PIXELS
idx2 = tf.linalg.matmul(m, tf.cast(
idx, dtype='float32')) # shape = (batch_size, 3, DIM ** 2)
idx2 = K.cast(idx2, dtype='int32') # shape = (batch_size, 3, DIM ** 2)
if PADDING == "REFLECT": #TPU compatible reflection padding, based on https://www.kaggle.com/psaikko/augmentations-with-reflect-padding
idx2 = tf.stack([DIM // 2 - idx2[:, 0, ], DIM // 2 + idx2[:, 1, ]],
axis=1)
# Identify out-of-bounds positions
bounds_mask = tf.math.logical_or(tf.math.less(idx2, 0),
tf.math.greater(idx2, DIM - 1))
# Compute mirrored positions
mirror_idxs = tf.math.subtract(DIM - 1,
tf.math.floormod(idx2, DIM - 1))
idx2 = tf.where(bounds_mask, mirror_idxs, idx2)
idx3 = tf.stack([idx2[:, 0, ], idx2[:, 1, ]], axis=1)
else:
#idx2 = K.clip(idx2,-DIM//2+XDIM+1,DIM//2)
idx3 = tf.stack([DIM // 2 - idx2[:, 0, ], DIM // 2 + idx2[:, 1, ]],
axis=1)
# shape = (batch_size, DIM ** 2, 3)
d = tf.gather_nd(images, tf.transpose(idx3, perm=[0, 2, 1]), batch_dims=1)
d = tf.reshape(d, (B, DIM, DIM, 3))
#d=tf.compat.v1.image.resize(d, [real_H*2,real_W*2])
d = tf.nn.avg_pool(d,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
data_format='NHWC')
# shape = (batch_size, DIM, DIM, 3)
images = tf.reshape(d, (B, DIM // 2, DIM // 2, 3))
#images=images[:,(real_H)//2:(3*real_H)//2,(real_W)//2:(3*real_W)//2,:]
return images
def _dataset_parser(value):
"""Parse data to a fixed dimension input image and learning targets.
Args:
value: A dictionary contains an image and groundtruth annotations.
Returns:
image: Image tensor that is preprocessed to have normalized value and
fixed dimension [image_height, image_width, 3]
cls_targets_dict: ordered dictionary with keys
[min_level, min_level+1, ..., max_level]. The values are tensor with
shape [height_l, width_l, num_anchors]. The height_l and width_l
represent the dimension of class logits at l-th level.
box_targets_dict: ordered dictionary with keys
[min_level, min_level+1, ..., max_level]. The values are tensor with
shape [height_l, width_l, num_anchors * 4]. The height_l and
width_l represent the dimension of bounding box regression output at
l-th level.
num_positives: Number of positive anchors in the image.
source_id: Source image id. Default value -1 if the source id is empty
in the groundtruth annotation.
image_scale: Scale of the processed image to the original image.
boxes: Groundtruth bounding box annotations. The box is represented in
[y1, x1, y2, x2] format. The tensor is padded with -1 to the fixed
dimension [self._max_instances_per_image, 4].
is_crowds: Groundtruth annotations to indicate if an annotation
represents a group of instances by value {0, 1}. The tensor is
padded with 0 to the fixed dimension [self._max_instances_per_image].
areas: Groundtruth areas annotations. The tensor is padded with -1
to the fixed dimension [self._max_instances_per_image].
classes: Groundtruth classes annotations. The tensor is padded with -1
to the fixed dimension [self._max_instances_per_image].
"""
with tf.name_scope('parser'):
data = example_decoder.decode(value)
source_id = data['source_id']
image = data['image']
boxes = data['groundtruth_boxes']
classes = data['groundtruth_classes']
classes = tf.reshape(tf.cast(classes, dtype=tf.float32),
[-1, 1])
areas = data['groundtruth_area']
is_crowds = data['groundtruth_is_crowd']
classes = tf.reshape(tf.cast(classes, dtype=tf.float32),
[-1, 1])
if params['skip_crowd_during_training'] and self._is_training:
indices = tf.where(
tf.logical_not(data['groundtruth_is_crowd']))
classes = tf.gather_nd(classes, indices)
boxes = tf.gather_nd(boxes, indices)
# NOTE: The autoaugment method works best when used alongside the
# standard horizontal flipping of images along with size jittering
# and normalization.
if params.get('autoaugment_policy',
None) and self._is_training:
from aug import autoaugment # pylint: disable=g-import-not-at-top
image, boxes = autoaugment.distort_image_with_autoaugment(
image, boxes, params['autoaugment_policy'],
params['use_augmix'], *params['augmix_params'])
input_processor = DetectionInputProcessor(
image, params['image_size'], boxes, classes)
input_processor.normalize_image()
if self._is_training and params['input_rand_hflip']:
input_processor.random_horizontal_flip()
if self._is_training:
input_processor.set_training_random_scale_factors(
params['train_scale_min'], params['train_scale_max'],
params.get('target_size', None))
else:
input_processor.set_scale_factors_to_output_size()
image = input_processor.resize_and_crop_image()
boxes, classes = input_processor.resize_and_crop_boxes()
# Assign anchors.
(cls_targets, box_targets,
num_positives) = anchor_labeler.label_anchors(boxes, classes)
source_id = tf.where(tf.equal(source_id, tf.constant('')),
'-1', source_id)
source_id = tf.string_to_number(source_id)
# Pad groundtruth data for evaluation.
image_scale = input_processor.image_scale_to_original
boxes *= image_scale
is_crowds = tf.cast(is_crowds, dtype=tf.float32)
boxes = pad_to_fixed_size(boxes, -1,
[self._max_instances_per_image, 4])
is_crowds = pad_to_fixed_size(
is_crowds, 0, [self._max_instances_per_image, 1])
areas = pad_to_fixed_size(areas, -1,
[self._max_instances_per_image, 1])
classes = pad_to_fixed_size(classes, -1,
[self._max_instances_per_image, 1])
return (image, cls_targets, box_targets, num_positives,
source_id, image_scale, boxes, is_crowds, areas,
classes)
def _build(self, x, presence=None):
batch_size, n_input_points = x.shape[:2].as_list()
# we don't know what order the initial points came in, so we need to create
# a big mixture of all votes for every input point
# [B, 1, n_votes, n_input_dims]
expanded_votes = tf.expand_dims(self._votes, 1)
expanded_scale = tf.expand_dims(tf.expand_dims(self._scales, 1), -1)
vote_component_pdf = self._get_pdf(expanded_votes, expanded_scale)
# [B, n_points, n_caps, n_votes, n_input_dims]
expanded_x = tf.expand_dims(x, 2)
vote_log_prob_per_dim = vote_component_pdf.log_prob(expanded_x)
# [B, n_points, n_votes]
vote_log_prob = tf.reduce_sum(vote_log_prob_per_dim, -1)
dummy_vote_log_prob = tf.zeros([batch_size, n_input_points, 1])
dummy_vote_log_prob -= 2. * tf.log(10.)
vote_log_prob = tf.concat([vote_log_prob, dummy_vote_log_prob], 2)
# [B, n_points, n_votes]
mixing_logits = math_ops.safe_log(self._vote_presence_prob)
dummy_logit = tf.zeros([batch_size, 1]) - 2. * tf.log(10.)
mixing_logits = tf.concat([mixing_logits, dummy_logit], 1)
mixing_log_prob = mixing_logits - tf.reduce_logsumexp(
mixing_logits, 1, keepdims=True)
expanded_mixing_logits = tf.expand_dims(mixing_log_prob, 1)
mixture_log_prob_per_component\
= tf.reduce_logsumexp(expanded_mixing_logits + vote_log_prob, 2)
if presence is not None:
presence = tf.cast(presence, tf.float32)
mixture_log_prob_per_component *= presence
mixture_log_prob_per_example\
= tf.reduce_sum(mixture_log_prob_per_component, 1)
mixture_log_prob_per_batch = tf.reduce_mean(
mixture_log_prob_per_example)
# [B, n_points, n_votes]
posterior_mixing_logits_per_point = expanded_mixing_logits + vote_log_prob
# [B, n_points]
winning_vote_idx = tf.argmax(
posterior_mixing_logits_per_point[:, :, :-1], 2)
batch_idx = tf.expand_dims(tf.range(batch_size, dtype=tf.int64), -1)
batch_idx = snt.TileByDim([1], [winning_vote_idx.shape[-1]])(batch_idx)
idx = tf.stack([batch_idx, winning_vote_idx], -1)
winning_vote = tf.gather_nd(self._votes, idx)
winning_pres = tf.gather_nd(self._vote_presence_prob, idx)
vote_presence = tf.greater(mixing_logits[:, :-1], mixing_logits[:,
-1:])
# the first four votes belong to the square
is_from_capsule = winning_vote_idx // self._n_votes
posterior_mixing_probs = tf.nn.softmax(
posterior_mixing_logits_per_point, -1)[Ellipsis, :-1]
assert winning_vote.shape == x.shape
return self.OutputTuple(
log_prob=mixture_log_prob_per_batch,
vote_presence=tf.cast(vote_presence, tf.float32),
winner=winning_vote,
winner_pres=winning_pres,
is_from_capsule=is_from_capsule,
mixing_logits=mixing_logits,
mixing_log_prob=mixing_log_prob,
# TODO(adamrk): this is broken
soft_winner=tf.zeros_like(winning_vote),
soft_winner_pres=tf.zeros_like(winning_pres),
posterior_mixing_probs=posterior_mixing_probs,
)
def build_distractors(distractor_examples, context):
"""Create inputs with distractors."""
CLS_ID = tf.constant([101], dtype=tf.int64) # pylint: disable=invalid-name
SEP_ID = tf.constant([102], dtype=tf.int64) # pylint: disable=invalid-name
bert_inputs = []
input_masks = []
segment_ids = []
# for each distractor
sample_size = int(
(FLAGS.num_choices - FLAGS.k_size) / (FLAGS.data_window_size - 1))
for example in distractor_examples:
# randomly sample 7
intermediate_examples_tensor = tf.reduce_sum(tf.abs(example), 1)
examples_zero_vector = tf.zeros(shape=(1, 1), dtype=tf.int64)
examples_bool_mask = tf.squeeze(
tf.not_equal(intermediate_examples_tensor, examples_zero_vector))
paragraph_len = tf.reduce_sum(tf.cast(examples_bool_mask, tf.int32))
indices = tf.range(0, limit=paragraph_len, dtype=tf.int32)
shuffled_indices = tf.random.shuffle(indices)[:sample_size]
# extend examples / targets
distractor_cand = example
distractor_cand_plus_one = distractor_cand[1:]
distractor_cand_plus_two = distractor_cand[2:]
# pad extensions
paddings_one = tf.constant([[0, 1], [0, 0]])
distractor_cand_plus_one = tf.pad(distractor_cand_plus_one, paddings_one)
paddings_two = tf.constant([[0, 2], [0, 0]])
distractor_cand_plus_two = tf.pad(distractor_cand_plus_two, paddings_two)
distractor_cand_ext = tf.concat(
[distractor_cand, distractor_cand_plus_one, distractor_cand_plus_two],
axis=1)
distractors = tf.gather(distractor_cand_ext, shuffled_indices)
for i in range(sample_size):
distractors_non_zero = tf.where(
tf.not_equal(distractors[i], tf.zeros_like(distractors[i])))
distractors_stripped = tf.gather_nd(distractors[i], distractors_non_zero)
if FLAGS.include_context:
segment_id = tf.concat([
tf.zeros_like(CLS_ID, dtype=tf.int64),
tf.zeros_like(context),
tf.zeros_like(SEP_ID, dtype=tf.int64),
tf.ones_like(distractors_stripped),
tf.ones_like(SEP_ID, dtype=tf.int64)
],
axis=0)
else:
segment_id = tf.concat([
tf.zeros_like(CLS_ID, dtype=tf.int64),
tf.zeros_like(distractors_stripped),
tf.zeros_like(SEP_ID, dtype=tf.int64)
],
axis=0)
segment_id = pad_and_cut(segment_id, FLAGS.max_seq_length)
segment_ids.append(segment_id)
if FLAGS.include_context:
new_input = tf.concat(
[CLS_ID, context, SEP_ID, distractors_stripped, SEP_ID], axis=0)
else:
new_input = tf.concat([CLS_ID, distractors_stripped, SEP_ID], axis=0)
input_mask = tf.ones_like(new_input)
input_mask = pad_and_cut(input_mask, FLAGS.max_seq_length)
input_masks.append(input_mask)
padded_new_input = pad_and_cut(new_input, FLAGS.max_seq_length)
bert_inputs.append(padded_new_input)
bert_inputs = tf.stack(bert_inputs, axis=0)
input_masks = tf.stack(input_masks, axis=0)
segment_ids = tf.stack(segment_ids, axis=0)
out = Outputs_And_Context(bert_inputs, input_masks, segment_ids, None, None)
return out
def _build(self, x, presence=None):
# x is [B, n_input_points, n_input_dims]
batch_size, n_input_points = x.shape[:2].as_list()
# votes and scale have shape [B, n_caps, n_input_points, n_input_dims|1]
# since scale is a per-caps scalar and we have one vote per capsule
vote_component_pdf = self._get_pdf(self._votes,
tf.expand_dims(self._scales, -1))
# expand along caps dimensions -> [B, 1, n_input_points, n_input_dims]
expanded_x = tf.expand_dims(x, 1)
vote_log_prob_per_dim = vote_component_pdf.log_prob(expanded_x)
# [B, n_caps, n_input_points]
vote_log_prob = tf.reduce_sum(vote_log_prob_per_dim, -1)
dummy_vote_log_prob = tf.zeros([batch_size, 1, n_input_points])
dummy_vote_log_prob -= 2. * tf.log(10.)
# [B, n_caps + 1, n_input_points]
vote_log_prob = tf.concat([vote_log_prob, dummy_vote_log_prob], 1)
# [B, n_caps, n_input_points]
mixing_logits = math_ops.safe_log(self._vote_presence_prob)
dummy_logit = tf.zeros([batch_size, 1, 1]) - 2. * tf.log(10.)
dummy_logit = snt.TileByDim([2], [n_input_points])(dummy_logit)
# [B, n_caps + 1, n_input_points]
mixing_logits = tf.concat([mixing_logits, dummy_logit], 1)
mixing_log_prob = mixing_logits - tf.reduce_logsumexp(
mixing_logits, 1, keepdims=True)
# [B, n_input_points]
mixture_log_prob_per_point = tf.reduce_logsumexp(
mixing_logits + vote_log_prob, 1)
if presence is not None:
presence = tf.cast(presence, tf.float32)
mixture_log_prob_per_point *= presence
# [B,]
mixture_log_prob_per_example\
= tf.reduce_sum(mixture_log_prob_per_point, 1)
# []
mixture_log_prob_per_batch = tf.reduce_mean(
mixture_log_prob_per_example)
# [B, n_caps + 1, n_input_points]
posterior_mixing_logits_per_point = mixing_logits + vote_log_prob
# [B, n_input_points]
winning_vote_idx = tf.argmax(posterior_mixing_logits_per_point[:, :-1],
1)
batch_idx = tf.expand_dims(tf.range(batch_size, dtype=tf.int64), 1)
batch_idx = snt.TileByDim([1], [n_input_points])(batch_idx)
point_idx = tf.expand_dims(tf.range(n_input_points, dtype=tf.int64), 0)
point_idx = snt.TileByDim([0], [batch_size])(point_idx)
idx = tf.stack([batch_idx, winning_vote_idx, point_idx], -1)
winning_vote = tf.gather_nd(self._votes, idx)
winning_pres = tf.gather_nd(self._vote_presence_prob, idx)
vote_presence = tf.greater(mixing_logits[:, :-1], mixing_logits[:,
-1:])
# the first four votes belong to the square
is_from_capsule = winning_vote_idx // self._n_votes
posterior_mixing_probs = tf.nn.softmax(
posterior_mixing_logits_per_point, 1)
dummy_vote = tf.get_variable('dummy_vote',
shape=self._votes[:1, :1].shape)
dummy_vote = snt.TileByDim([0], [batch_size])(dummy_vote)
dummy_pres = tf.zeros([batch_size, 1, n_input_points])
votes = tf.concat((self._votes, dummy_vote), 1)
pres = tf.concat([self._vote_presence_prob, dummy_pres], 1)
soft_winner = tf.reduce_sum(
tf.expand_dims(posterior_mixing_probs, -1) * votes, 1)
soft_winner_pres = tf.reduce_sum(posterior_mixing_probs * pres, 1)
posterior_mixing_probs = tf.transpose(posterior_mixing_probs[:, :-1],
(0, 2, 1))
assert winning_vote.shape == x.shape
return self.OutputTuple(
log_prob=mixture_log_prob_per_batch,
vote_presence=tf.cast(vote_presence, tf.float32),
winner=winning_vote,
winner_pres=winning_pres,
soft_winner=soft_winner,
soft_winner_pres=soft_winner_pres,
posterior_mixing_probs=posterior_mixing_probs,
is_from_capsule=is_from_capsule,
mixing_logits=mixing_logits,
mixing_log_prob=mixing_log_prob,
)
def body_cross_entropy(itr, cond_terminate, sample_mean,
sample_covariance_diag, top_k_value,
top_k_action_samples):
"""Function for cross entropy search of actions."""
del top_k_action_samples
top_k_value_prev = top_k_value
batch_sample_mean = tf.reshape(
tf.tile(sample_mean, [1, num_samples]),
[self._dynamic_batch_size * num_samples, self.action_dim])
batch_sample_covariance_diag = tf.reshape(
tf.tile(sample_covariance_diag, [1, num_samples]),
[self._dynamic_batch_size * num_samples, self.action_dim])
action_samples = self._action_projection(
batch_sample_mean +
batch_sample_covariance_diag * tf.cast(random_sampler.sample(
sample_shape=[self._dynamic_batch_size * num_samples]),
dtype=tf.float32))
state_samples = tf.reshape(
tf.tile(self._state_tensor, [1, num_samples]),
[self._dynamic_batch_size * num_samples, self.state_dim])
action_samples = tf.reshape(
action_samples,
[self._dynamic_batch_size * num_samples, self.action_dim])
values = tf.reshape(
self._build_q_function_net(state_samples, action_samples),
[self._dynamic_batch_size, num_samples])
# everything is in batch mode
top_k_index = tf.argsort(values, axis=1,
direction="DESCENDING")[:, 0:top_k_num]
top_k_index_1d = tf.reshape(
top_k_index, [self._dynamic_batch_size * top_k_num, 1])
counter_tensor_1d = tf.reshape(
tf.tile(
tf.reshape(tf.range(self._dynamic_batch_size),
[self._dynamic_batch_size, 1]), [1, top_k_num]),
[self._dynamic_batch_size * top_k_num, 1])
top_k_index_2d = tf.concat([counter_tensor_1d, top_k_index_1d],
axis=1)
action_samples = tf.reshape(
action_samples,
[self._dynamic_batch_size, num_samples, self.action_dim])
top_k_action_samples = tf.gather_nd(action_samples, top_k_index_2d)
top_k_action_samples = tf.reshape(
top_k_action_samples,
[self._dynamic_batch_size, top_k_num, self.action_dim])
top_k_values = tf.gather_nd(values, top_k_index_2d)
top_k_values = tf.reshape(top_k_values,
[self._dynamic_batch_size, top_k_num])
# it's a batch_size x 1 tensor
top_k_value = tf.reshape(tf.reduce_mean(top_k_values, axis=1),
[self._dynamic_batch_size, 1])
sample_mean = tf.reduce_mean(top_k_action_samples, axis=1)
sample_covariance_diag = tf.math.reduce_variance(
top_k_action_samples, axis=1)
itr = itr + 1
cond_terminate = tf.less_equal(
tf.reduce_mean(tf.math.abs(top_k_value - top_k_value_prev)),
self._tolerance_tensor)
return itr, cond_terminate, sample_mean, sample_covariance_diag, \
top_k_value, top_k_action_samples
def flatten_padded_sequences():
indices = tf.where(tf.sequence_mask(lengths))
return tf.gather_nd(maybe_padded_sequences, indices)
def box_matching(boxes, gt_boxes, gt_classes, gt_attributes):
"""Match boxes to groundtruth boxes.
Given the proposal boxes and the groundtruth boxes, classes and attributes,
perform the groundtruth matching by taking the argmax of the IoU between boxes
and groundtruth boxes.
Args:
boxes: a tensor of shape of [batch_size, N, 4] representing the box
coordiantes to be matched to groundtruth boxes.
gt_boxes: a tensor of shape of [batch_size, MAX_INSTANCES, 4] representing
the groundtruth box coordinates. It is padded with -1s to indicate the
invalid boxes.
gt_classes: [batch_size, MAX_INSTANCES] representing the groundtruth box
classes. It is padded with -1s to indicate the invalid classes.
gt_attributes: [batch_size, MAX_NUM_INSTANCES, num_attributes] representing
the groundtruth attributes. It is padded with -1s to indicate the invalid
attributes.
Returns:
matched_gt_boxes: a tensor of shape of [batch_size, N, 4], representing
the matched groundtruth box coordinates for each input box. If the box
does not overlap with any groundtruth boxes, the matched boxes of it
will be set to all 0s.
matched_gt_classes: a tensor of shape of [batch_size, N], representing
the matched groundtruth classes for each input box. If the box does not
overlap with any groundtruth boxes, the matched box classes of it will
be set to 0, which corresponds to the background class.
matched_gt_attributes: a tensor of shape of [batch_size, N,
num_attributes], representing the matched groundtruth attributes for each
input box. If the box does not overlap with any groundtruth boxes, the
matched box attributes of it will be set to all 0s.
matched_gt_indices: a tensor of shape of [batch_size, N], representing
the indices of the matched groundtruth boxes in the original gt_boxes
tensor. If the box does not overlap with any groundtruth boxes, the
index of the matched groundtruth will be set to -1.
matched_iou: a tensor of shape of [batch_size, N], representing the IoU
between the box and its matched groundtruth box. The matched IoU is the
maximum IoU of the box and all the groundtruth boxes.
iou: a tensor of shape of [batch_size, N, K], representing the IoU matrix
between boxes and the groundtruth boxes. The IoU between a box and the
invalid groundtruth boxes whose coordinates are [-1, -1, -1, -1] is -1.
"""
# Compute IoU between boxes and gt_boxes.
# iou <- [batch_size, N, K]
iou = box_utils.bbox_overlap(boxes, gt_boxes)
# max_iou <- [batch_size, N]
# 0.0 -> no match to gt, or -1.0 match to no gt
matched_iou = tf.reduce_max(iou, axis=-1)
# background_box_mask <- bool, [batch_size, N]
background_box_mask = tf.less_equal(matched_iou, 0.0)
argmax_iou_indices = tf.argmax(iou, axis=-1, output_type=tf.int32)
argmax_iou_indices_shape = tf.shape(argmax_iou_indices)
batch_indices = (
tf.expand_dims(tf.range(argmax_iou_indices_shape[0]), axis=-1) *
tf.ones([1, argmax_iou_indices_shape[-1]], dtype=tf.int32))
gather_nd_indices = tf.stack([batch_indices, argmax_iou_indices], axis=-1)
matched_gt_boxes = tf.gather_nd(gt_boxes, gather_nd_indices)
matched_gt_boxes = tf.where(
tf.tile(tf.expand_dims(background_box_mask, axis=-1), [1, 1, 4]),
tf.zeros_like(matched_gt_boxes, dtype=tf.float32), matched_gt_boxes)
matched_gt_classes = tf.gather_nd(gt_classes, gather_nd_indices)
matched_gt_classes = tf.where(background_box_mask,
tf.zeros_like(matched_gt_classes),
matched_gt_classes)
_, _, num_attributes = gt_attributes.get_shape().as_list()
matched_gt_attributes = tf.gather_nd(gt_attributes, gather_nd_indices)
matched_gt_attributes = tf.where(
tf.tile(tf.expand_dims(background_box_mask, axis=-1),
[1, 1, num_attributes]),
tf.zeros_like(matched_gt_attributes, dtype=tf.float32),
matched_gt_attributes)
matched_gt_indices = tf.where(background_box_mask,
-tf.ones_like(argmax_iou_indices),
argmax_iou_indices)
return (matched_gt_boxes, matched_gt_classes, matched_gt_attributes,
matched_gt_indices, matched_iou, iou)
def __init__(self, lr, batch_size, dimension, util_train, util_test, campaign, reg_lambda, sigma):
# hyperparameters
self.lr = lr
self.batch_size = batch_size
self.util_train = util_train
self.util_test = util_test
self.reg_lambda = reg_lambda
self.sigma = sigma
self.emb_size = 20
self.train_data_amt = util_train.get_data_amt()
self.test_data_amt = util_test.get_data_amt()
# output dir
model_name = "{}_{}_{}_{}".format(self.lr, self.reg_lambda, self.batch_size, self.sigma)
self.output_dir = "output/deephit/{}/{}/".format(campaign, model_name)
if not os.path.exists(self.output_dir):
os.makedirs(self.output_dir)
# reset graph
tf.reset_default_graph()
# field params
self.field_sizes = self.util_train.feat_sizes
self.field_num = len(self.field_sizes)
# placeholders
self.X = [tf.sparse_placeholder(tf.float64) for i in range(0, self.field_num)]
self.z = tf.placeholder(tf.float64)
self.b = tf.placeholder(tf.float64)
self.y = tf.placeholder(tf.float64)
# embedding layer
self.var_map = {}
# for truncated
self.var_map['embed_0'] = tf.Variable(
tf.truncated_normal([self.field_sizes[0], 1], dtype=tf.float64))
for i in range(1, self.field_num):
self.var_map['embed_%d' % i] = tf.Variable(
tf.truncated_normal([self.field_sizes[i], self.emb_size], dtype=tf.float64))
# after embedding
w0 = [self.var_map['embed_%d' % i] for i in range(self.field_num)]
self.dense_input = tf.concat([tf.sparse_tensor_dense_matmul(self.X[i], w0[i]) for i in range(self.field_num)], 1)
# shared network
self.hidden1 = tf.Variable(initial_value=tf.truncated_normal(shape=[(self.field_num - 1) * self.emb_size + 1, HIDDEN_SIZE1], dtype=tf.float64), name='h1')
self.out1 = tf.Variable(initial_value=tf.truncated_normal(shape=[HIDDEN_SIZE1, OUT_SIZE1], dtype=tf.float64), name='o1')
self.hidden2 = tf.Variable(initial_value=tf.truncated_normal(shape=[OUT_SIZE1, HIDDEN_SIZE2], dtype=tf.float64), name='h2')
self.out2 = tf.Variable(initial_value=tf.truncated_normal(shape=[HIDDEN_SIZE2, OUT_SIZE2], dtype=tf.float64), name='o2')
# cause-specific network
self.hidden1_val = tf.nn.relu(tf.matmul(self.dense_input, self.hidden1))
self.out1_val = tf.sigmoid(tf.matmul(self.hidden1_val, self.out1))
self.hidden2_val = tf.nn.relu(tf.matmul(self.out1_val, self.hidden2))
self.out2_val = tf.sigmoid(tf.matmul(self.hidden2_val, self.out2))
# p_z and w_b
self.p = tf.nn.softmax(self.out2_val)
self.w = tf.cumsum(self.p, exclusive=True, axis = 1)
idx_z = tf.stack([tf.reshape(tf.range(tf.shape(self.z)[0]), (-1,1)), tf.cast(self.z - 1, tf.int32)], axis=-1)
idx_b = tf.stack([tf.reshape(tf.range(tf.shape(self.b)[0]), (-1,1)), tf.cast(self.b - 1, tf.int32)], axis=-1)
self.pz = tf.gather_nd(self.p, idx_z)
self.wb = tf.gather_nd(self.w, idx_b)
self.wz = tf.gather_nd(self.w, idx_z)
# loss and train step
self.loss1 = -tf.reduce_sum(tf.log(tf.clip_by_value(self.pz, 1e-8, 1.0)) * self.y)
self.loss2 = -tf.reduce_sum(tf.log(tf.clip_by_value(1 - self.wb, 1e-8, 1.0)) * (1 - self.y))
self.reg_loss = tf.nn.l2_loss(self.hidden1[1:,]) + tf.nn.l2_loss(self.hidden2[1:,]) + \
tf.nn.l2_loss(self.out1[1:,]) + tf.nn.l2_loss(self.out2[1:,])
# get ranking loss
self.w_of_pair = tf.transpose(tf.nn.embedding_lookup(tf.transpose(self.w), tf.cast(self.z[:,0] - 1, tf.int32)))
self.w_of_self = tf.reshape(tf.tile(tf.reshape(self.wz, (self.batch_size, )), [self.batch_size]), (self.batch_size, self.batch_size))
self.win_label = tf.reshape(tf.tile(tf.reshape(self.y, (self.batch_size, )), [self.batch_size]), (self.batch_size, self.batch_size))
self.delta = self.w_of_self - self.w_of_pair
self.candidate = tf.exp(-self.delta / self.sigma)
self.rank_loss = tf.reduce_sum(tf.matrix_band_part(self.candidate, -1, 0) * self.win_label)
self.loss = self.loss1 + self.loss2 + self.reg_lambda * self.reg_loss + self.rank_loss
self.optimizer = tf.train.GradientDescentOptimizer(self.lr)
self.train_step = self.optimizer.minimize(self.loss)
# session initialization
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
self.sess = tf.Session(config=config)
tf.global_variables_initializer().run(session=self.sess)
def _update_mask(self, weights, threshold, gradients): # pylint: disable=unused-argument
"""Updates the mask for a given weight tensor.
This functions first computes the cdf of the weight tensor, and estimates
the threshold value such that 'desired_sparsity' fraction of weights
have magnitude less than the threshold.
Args:
weights: The weight tensor that needs to be masked.
threshold: The current threshold value. The function will compute a new
threshold and return the exponential moving average using the current
value of threshold
gradients: The gradient tensor that is used for salience calculation.
Returns:
new_threshold: The new value of the threshold based on weights, and
sparsity at the current global_step
new_mask: A numpy array of the same size and shape as weights containing
0 or 1 to indicate which of the values in weights falls below
the threshold
Raises:
ValueError: if sparsity is not defined
"""
if self._sparsity is None:
raise ValueError('Sparsity variable undefined')
sparsity = self._get_sparsity(weights.op.name)
with tf.name_scope(weights.op.name + '_pruning_ops'):
tf.logging.info('Applying option %s pruning',
self._spec.prune_option)
if self._spec.prune_option == 'weight':
abs_weights = tf.abs(weights)
elif self._spec.prune_option in ('first_order_gradient',
'second_order_gradient'):
if gradients is None:
raise ValueError('gradient tensor cannot be None.')
# gradient variable stores absolute value already
abs_weights = tf.multiply(tf.abs(weights), gradients)
else:
raise ValueError('undefined option')
k = tf.cast(
tf.round(
tf.cast(tf.size(abs_weights), tf.float32) *
(1 - sparsity)), tf.int32)
# Generate a random shuffling of the weights s.t. the tie-breaker on
# weight magnitude is random uniform.
shuffling = tf.random_shuffle(tf.range(tf.size(abs_weights)))
shuffling = tf.reshape(shuffling, [-1, 1])
# Flatten the weights and scatter the values randomly.
abs_weights = tf.reshape(abs_weights, [-1])
abs_weights = tf.scatter_nd(shuffling, abs_weights,
tf.shape(abs_weights))
# Sort the entire array
_, indices = tf.nn.top_k(abs_weights, k=tf.size(abs_weights))
# `k` is how many non-zero weights we're going to have. Create a new
# mask where the first `k` elements are set to one and all others are
# set to zero.
mask_staging = tf.range(tf.size(abs_weights))
mask_staging = tf.cast(tf.less(mask_staging, k), tf.float32)
# Scatter the mask back into the proper positions for the weight matrix.
indices = tf.reshape(indices, [-1, 1])
new_mask = tf.scatter_nd(indices, mask_staging,
tf.shape(mask_staging))
# Un-shuffle the newly created mask.
new_mask = tf.reshape(tf.gather_nd(new_mask, shuffling),
tf.shape(weights))
return tf.constant(0, tf.float32), new_mask
def _build_train_op(self):
"""Builds a training op.
Returns:
train_op: An op performing one step of training from replay data.
"""
target_distribution = tf.stop_gradient(
self._build_target_distribution())
# size of indices: batch_size x 1.
indices = tf.range(tf.shape(self._replay_net_outputs.logits)[0])[:, None]
# size of reshaped_actions: batch_size x 2.
reshaped_actions = tf.concat([indices, self._replay.actions[:, None]], 1)
# For each element of the batch, fetch the logits for its selected action.
chosen_action_logits = tf.gather_nd(self._replay_net_outputs.logits,
reshaped_actions)
loss = tf.nn.softmax_cross_entropy_with_logits(
labels=target_distribution,
logits=chosen_action_logits)
# Record returns encountered in the sampled training batches.
returns = self._replay.transition['return']
statistics_summaries('returns', returns)
train_counts = self._replay.transition['train_counts']
statistics_summaries('train_counts', train_counts)
steps_until_first_train = self._replay.transition['steps_until_first_train']
statistics_summaries('steps_until_first_train', steps_until_first_train)
age = self._replay.transition['age']
statistics_summaries('age', age)
if self._replay_scheme == 'prioritized':
# The original prioritized experience replay uses a linear exponent
# schedule 0.4 -> 1.0. Comparing the schedule to a fixed exponent of 0.5
# on 5 games (Asterix, Pong, Q*Bert, Seaquest, Space Invaders) suggested
# a fixed exponent actually performs better, except on Pong.
probs = self._replay.transition['sampling_probabilities']
beta = self._beta_exponent
tf.summary.histogram('probs', probs)
loss_weights = 1.0 / tf.math.pow(probs + 1e-10, beta)
tf.summary.histogram('loss_weights', loss_weights)
loss_weights /= tf.reduce_max(loss_weights)
# Rainbow and prioritized replay are parametrized by an exponent alpha,
# but in both cases it is set to 0.5 - for simplicity's sake we leave it
# as is here, using the more direct tf.sqrt(). Taking the square root
# "makes sense", as we are dealing with a squared loss.
# Add a small nonzero value to the loss to avoid 0 priority items. While
# technically this may be okay, setting all items to 0 priority will cause
# troubles, and also result in 1.0 / 0.0 = NaN correction terms.
alpha = self._alpha_exponent
update_priorities_op = self._replay.tf_set_priority(
self._replay.indices, tf.math.pow(loss + 1e-10, alpha))
# Weight the loss by the inverse priorities.
loss = loss_weights * loss
else:
update_priorities_op = tf.no_op()
update_train_counts_op = self._replay.tf_update_train_counts(
self._replay.indices)
with tf.control_dependencies([update_priorities_op,
update_train_counts_op]):
if self.summary_writer is not None:
with tf.variable_scope('Losses'):
tf.summary.scalar('CrossEntropyLoss', tf.reduce_mean(loss))
# Schaul et al. reports a slightly different rule, where 1/N is also
# exponentiated by beta. Not doing so seems more reasonable, and did not
# impact performance in our experiments.
return self.optimizer.minimize(tf.reduce_mean(loss)), loss
