def transform_targets_for_output(y_true, grid_size, anchor_idxs, classes):
"""
y_true here is y_train append a column with anchor_idx
which indicates, every object(box) anchored to which anchor
then you can using anchor to trace location and regression the precise offset to
the anchor
"""
# y_true: (N, boxes, (x1, y1, x2, y2, class, best_anchor))
N = tf.shape(y_true)[0]
# y_true_out: (N, grid, grid, anchors, [x, y, w, h, obj, class])
y_true_out = tf.zeros(
(N, grid_size, grid_size, tf.shape(anchor_idxs)[0], 6))
anchor_idxs = tf.cast(anchor_idxs, tf.int32)
indexes = tf.TensorArray(tf.int32, 1, dynamic_size=True)
updates = tf.TensorArray(tf.float32, 1, dynamic_size=True)
idx = 0
for i in tf.range(N):
for j in tf.range(tf.shape(y_true)[1]):
if tf.equal(y_true[i][j][2], 0):
# 0 is backgrond or width is 0?
continue
anchor_eq = tf.equal(anchor_idxs, tf.cast(y_true[i][j][5],
tf.int32))
if tf.reduce_any(anchor_eq):
box = y_true[i][j][0:4]
box_xy = (y_true[i][j][0:2] + y_true[i][j][2:4]) / 2
anchor_idx = tf.cast(tf.where(anchor_eq), tf.int32)
grid_xy = tf.cast(box_xy // (1 / grid_size), tf.int32)
# grid[y][x][anchor] = (tx, ty, bw, bh, obj, class)
indexes = indexes.write(
idx, [i, grid_xy[1], grid_xy[0], anchor_idx[0][0]])
updates = updates.write(
idx, [box[0], box[1], box[2], box[3], 1, y_true[i][j][4]])
idx += 1
# tf.print(indexes.stack())
# tf.print(updates.stack())
return tf.tensor_scatter_nd_update(y_true_out, indexes.stack(),
updates.stack())
def segment_top_k(x, i, ratio):
"""
Returns indices to get the top K values in x segment-wise, according to
the segments defined in I. K is not fixed, but it is defined as a ratio of
the number of elements in each segment.
:param x: a rank 1 Tensor;
:param i: a rank 1 Tensor with segment IDs for x;
:param ratio: float, ratio of elements to keep for each segment;
:return: a rank 1 Tensor containing the indices to get the top K values of
each segment in x.
"""
i = tf.cast(i, tf.int32)
n = tf.shape(i)[0]
n_nodes = tf.math.segment_sum(tf.ones_like(i), i)
batch_size = tf.shape(n_nodes)[0]
n_nodes_max = tf.reduce_max(n_nodes)
cumulative_n_nodes = tf.concat(
(tf.zeros(1, dtype=n_nodes.dtype), tf.cumsum(n_nodes)[:-1]), 0
)
index = tf.range(n)
index = (index - tf.gather(cumulative_n_nodes, i)) + (i * n_nodes_max)
dense_x = tf.zeros(batch_size * n_nodes_max, dtype=x.dtype) - 1e20
dense_x = tf.tensor_scatter_nd_update(dense_x, index[:, None], x)
dense_x = tf.reshape(dense_x, (batch_size, n_nodes_max))
perm = tf.argsort(dense_x, direction="DESCENDING")
perm = perm + cumulative_n_nodes[:, None]
perm = tf.reshape(perm, (-1,))
k = tf.cast(tf.math.ceil(ratio * tf.cast(n_nodes, tf.float32)), i.dtype)
# This costs more memory
# to_rep = tf.tile(tf.constant([1., 0.]), (batch_size,))
# rep_times = tf.reshape(tf.concat((k[:, None], (n_nodes_max - k)[:, None]), -1), (-1,))
# mask = ops.repeat(to_rep, rep_times)
# perm = tf.boolean_mask(perm, mask)
# This is slower
r_range = tf.ragged.range(k).flat_values
r_delta = ops.repeat(tf.range(batch_size) * n_nodes_max, k)
mask = r_range + r_delta
perm = tf.gather(perm, mask)
return perm
def match_boxes(gt_labels, def_boxes, threshold=0.5):
"""
Return the new ground truth, corresponding to labeled default boxes
Parameters
----------
gt_labels: Tensor of real bounding boxes of the image(s), with labels, in the format:
[NUM_GT_BOXES * [x_min, y_min, x_max, y_max | label]]
def_boxes: Tensor of default boxes of an image, in the format:
[NUM_DEF_BOXES * [x_min, y_min, x_max, y_max]]
threshold: positive threshold
Return
------
ground_truth: labeled default boxes
"""
# Handle missing labels
if len(gt_labels) == 0:
def_labels = tf.zeros(shape=(len(def_boxes), ), dtype=tf.float32)
return def_boxes, def_labels
# Process ground truth
gt_boxes = gt_labels[..., :-1]
labels = gt_labels[..., -1]
# Compute IoU values
jaccard = jaccard_overlap(def_boxes, gt_boxes)
max_gt_jaccard = tf.reduce_max(jaccard, axis=1)
max_def_jaccard = tf.reduce_max(jaccard, axis=0)
max_gt_indices = tf.argmax(jaccard, axis=1)
max_def_indices = tf.argmax(jaccard, axis=0)
# Ensure best IoU
max_gt_jaccard = tf.tensor_scatter_nd_update(
max_gt_jaccard, tf.expand_dims(max_def_indices, 1),
tf.ones_like(max_def_indices, dtype=tf.float32))
# Match boxes with IoU > threshold
positive_mask = tf.where(max_gt_jaccard >= threshold, 1., 0.)
matched_labels = tf.gather(labels, max_gt_indices) * positive_mask
matched_boxes = tf.gather(gt_boxes, max_gt_indices)
encoded_boxes = encode_boxes(def_boxes, matched_boxes)
return encoded_boxes, matched_labels
def mask(self):
"""mask helper function for initializing grid mask of required size."""
mask_w = mask_h = int(
(self.gridmask_size_ratio + 1) * max(self.h, self.w))
mask = tf.zeros(shape=[mask_h, mask_w], dtype=tf.int32)
gridblock = tf.random.uniform(
shape=[],
minval=int(min(self.h * 0.5, self.w * 0.3)),
maxval=int(max(self.h * 0.5, self.w * 0.3)),
dtype=tf.int32,
)
if self.ratio == 1:
length = tf.random.uniform(shape=[],
minval=1,
maxval=gridblock,
dtype=tf.int32)
else:
length = tf.cast(
tf.math.minimum(
tf.math.maximum(
int(tf.cast(gridblock, tf.float32) * self.ratio + 0.5),
1,
),
gridblock - 1,
),
tf.int32,
)
for _ in range(2):
start_w = tf.random.uniform(shape=[],
minval=0,
maxval=gridblock,
dtype=tf.int32)
for i in range(mask_w // gridblock):
start = gridblock * i + start_w
end = tf.math.minimum(start + length, mask_w)
indices = tf.reshape(tf.range(start, end), [end - start, 1])
updates = (
tf.ones(shape=[end - start, mask_w], dtype=tf.int32) *
self.fill)
mask = tf.tensor_scatter_nd_update(mask, indices, updates)
mask = tf.transpose(mask)
return mask
def _sparse_scatter(
operand,
scatter_indices,
updates,
update_jaxpr,
update_consts,
dimension_numbers,
indices_are_sorted: bool,
unique_indices: bool,
mode,
_in_avals: Sequence[core.ShapedArray],
_out_aval: core.ShapedArray):
"""
Implementation of scatter specialised to indexing from the
front axes. This covers unique indices and non-unique indices
of single depth.
Note on unique indices: `tf.tensor_scatter_nd_update` interprets
indices thusly: every axis except the final one encodes a batch
dimension, the final axis encoding the actual indices to scatter in to.
It enforces, at least one, batch dimension so we add an empty
dimension to indices and updates if lacking.
Note on non-unique indices: There is no tf op for non single depth
indexing. But if indexing is single depth, this can be viewed as a
segment op.
"""
# Infer unique indices from lack of batch dimension
unique_indices = unique_indices or (len(scatter_indices.shape) == 1)
if unique_indices:
suboperand = tf.gather_nd(operand, scatter_indices)
updated_suboperand = update_op(suboperand, updates)
# add a batch dim if none exist
if len(scatter_indices.shape) == 1:
scatter_indices = scatter_indices[None]
updated_suboperand = updated_suboperand[None]
y = tf.tensor_scatter_nd_update(operand, scatter_indices, updated_suboperand)
else:
if (scatter_indices.shape[-1] == 1) and (unsorted_segment_op != None):
# If only indexing into the first dimension, it's a segment op
operand_update = unsorted_segment_op(updates, tf.squeeze(scatter_indices, -1), operand.shape[0])
y = update_op(operand, operand_update)
else:
raise error("Scatter supports unique indices. Scatter also supports non-unique indices with indexing into only one dimension for (add, mul, min, max)")
return y
def shiftGaugeField(self, gaugeField, cpt, sign):
gaugeFieldShifted = tf.roll(gaugeField, -sign, cpt)
pauliMatNum = self.boundaryConditions[cpt]
if pauliMatNum == 0:
return gaugeFieldShifted
latShape = tf.shape(gaugeField)[0:3]
indices = FieldTools.boundaryIndices(latShape, cpt, sign)
updates = tf.gather_nd(gaugeFieldShifted, indices)
updates = FieldTools.pauliMatrix(pauliMatNum) @\
updates @ FieldTools.pauliMatrix(pauliMatNum)
gaugeFieldShifted = tf.tensor_scatter_nd_update(
gaugeFieldShifted, indices, updates)
return gaugeFieldShifted
def sparse_fill_empty_rows(sp_input, default_value, name=None):
# Op implemented in tflite mode, so don't go through complicated logic below.
# if not cuda_ops_only():
# return tf.sparse_fill_empty_rows(sp_input, default_value, name=name)
d = sparse_tensor_to_dense(sp_input, default_value=default_value)
ds = sp_input.dense_shape
seq_len = sequence_length_from_sparse_tensor(sp_input)
empty_row_indicator = tf.equal(seq_len, 0)
mask = sparse_tensor_to_dense(tf.SparseTensor(sp_input.indices,
tf.fill(tf.shape(sp_input.values), True), ds))
row_idx = tf.where(tf.logical_and(empty_row_indicator,
tf.greater(tf.size(d), 0))) # empty input edge case.
fill_indices = tf.concat([row_idx, tf.zeros([tf.size(row_idx),
tf.size(ds) - 1], dtype=tf.int64)], axis=1)
mask = tf.tensor_scatter_nd_update(
mask, fill_indices, tf.fill(tf.shape(row_idx)[:1], True))
values, indices = tf.boolean_mask(d, mask), tf.where(mask)
return tf.SparseTensor(indices, values, ds), empty_row_indicator
def call(self, inputs):
Q_t, ard, Q_tar = inputs
Q_t1_max = tf.math.reduce_max(Q_tar, axis=1)
reward = tf.cast(ard[:, 1], tf.float32)
action = tf.cast(ard[:, 0], tf.int32)
done = tf.cast(ard[:, 2], tf.int32)
continue_candidates = tf.math.add(
tf.math.multiply(self.gamma, Q_t1_max), reward)
both_candidates = tf.stack([continue_candidates, reward],
axis=1)
new_Qta = tf.gather_nd(both_candidates,
tf.expand_dims(done, axis=1),
batch_dims=1)
size = tf.shape(ard)[0]
index = tf.range(0, size)
indices = tf.stack([index, action], axis=1)
new_Q = tf.tensor_scatter_nd_update(Q_t, indices, new_Qta)
return tf.stop_gradient(new_Q)
def shiftCovDeriv(self, covDeriv, cpt, sign):
covDerivShifted = tf.roll(covDeriv, -sign, cpt)
if cpt != 0:
return covDerivShifted
indices = FieldTools.boundaryIndices(self.latShape, cpt, sign)
if sign == -1:
updates = tf.zeros(tf.concat([tf.shape(indices)[0:-1], [2, 2]], 0),
dtype=tf.complex128)
else:
updates = tf.gather_nd(covDeriv, indices)
covDerivShifted = tf.tensor_scatter_nd_update(covDerivShifted, indices,
updates)
return covDerivShifted
def find_neighbouring_indices_along_axis(grid_sizes,
quadrature_coords_along_axis, dim):
dim_size = grid_sizes[dim]
dx = tf.cast(2 / (dim_size - 1), tf.keras.backend.floatx())
lower_idx = tf.cast(tf.math.floor((quadrature_coords_along_axis + 1) / dx),
tf.int32)
upper_idx = tf.cast(tf.math.ceil((quadrature_coords_along_axis + 1) / dx),
tf.int32)
equal_indices = (upper_idx == lower_idx)
if tf.reduce_any(
equal_indices
): #sometimes a quad pt may coincide with the position of a gridpoint. adjust upper idx for that possibility
new_upper_indices = tf.gather_nd(upper_idx,
tf.where(equal_indices)) + 1
upper_idx = tf.tensor_scatter_nd_update(upper_idx,
tf.where(equal_indices),
new_upper_indices)
return tf.stack([lower_idx, upper_idx], 1)
def pairwise_update_in_mol(h_v, mol_idx):
# (n_atoms, )
this_mol_has_atom = atom_in_mol[:, mol_idx]
# (n_atoms_in_this_mol, )
this_mol_atom_idxs = tf.where(this_mol_has_atom)
# (n_atoms_in_this_mol, d_v)
h_v_this_mol = tf.boolean_mask(h_v, this_mol_has_atom)
# (n_atoms_in_this_mol, d_v)
h_v_this_mol = self.pairwise_update(h_v_this_mol)
# (n_atoms, )
h_v = tf.tensor_scatter_nd_update(h_v, this_mol_atom_idxs,
h_v_this_mol)
return h_v, mol_idx + 1
def false_fn():
scatter_index = count_non_blank(y_hat_prediction, blank=self.text_featurizer.blank)
updated_prediction = tf.tensor_scatter_nd_update(
y_hat_prediction,
indices=tf.reshape(scatter_index, [1, 1]),
updates=tf.expand_dims(pred, axis=-1)
)
return (
B.score,
B.indices,
B.prediction,
B.states,
A.score.write(A_i, new_score),
A.indices.write(A_i, pred),
A.prediction.write(A_i, updated_prediction),
A.states.write(A_i, new_states),
A_i + 1
)
def kdtree_triplet_loss(self, y_true, y_pred, margin=1.0):
"""Computes the triplet loss with semi-hard negative mining.
MODIFIED FROM tensorflow_addons.losses.triplet_semihard_loss
Args:
y_true: 1-D integer `Tensor` with shape [batch_size] of
multiclass integer labels.
y_pred: 2-D float `Tensor` of embedding vectors. Embeddings should
be l2 normalized.
margin: Float, margin term in the loss definition.
"""
labels, embeddings = y_true, y_pred
# Reshape [batch_size] label tensor to a [batch_size, 1] label tensor.
lshape = tf.shape(labels)
print('step1')
assert lshape.shape == 1 or lshape.shape == (2, )
print('step2')
labels = tf.reshape(labels, [-1])
print('step3')
pos_dists, neg_dists = tf.py_function(self.get_dists,
inp=[labels, embeddings],
Tout=[tf.float64, tf.float64])
print('step4')
print(pos_dists, neg_dists)
pos_diffs = tf.minimum(pos_dists - neg_dists, [0])
print('step5')
print(pos_diffs)
zero_idx = tf.cast(
tf.where(pos_diffs == 0, -tf.ones_like(pos_diffs), pos_diffs),
tf.int64)
print(zero_idx)
print(pos_dists[zero_idx])
norm_diffs = tf.tensor_scatter_nd_update(pos_diffs, zero_idx,
pos_dists[zero_idx])
print(norm_diffs)
# pos_diffs[[1] == 0] = pos_dists[[1] == 0] / 2
print('step6')
triplet_loss = tf.cast(tf.reduce_mean(norm_diffs), tf.float32)
print('step7')
print(triplet_loss)
return triplet_loss
def patch_application(images: tf.Tensor,
patches: tf.Tensor,
randomness: Union[np.array, tf.Tensor] = None):
"""Applies patches to images at random locations.
Args:
images: A batch of images of shape [B, H, W, 3].
patches: A batch of patches of shape [B, H_patch, W_patch, 3].
randomness: Optional seeds for randomness. Shape is [B, ]
Returns:
Images with patches applied at random locations. Shape is [B, H, W, 3].
"""
image_height, image_width = images.shape[1:3]
patch_height, patch_width = patches.shape[1:3]
height_to_pad = int(image_height - patch_height)
width_to_pad = int(image_width - patch_width)
# Generate offsets for random translation
if randomness is None:
height_offsets = tf.random.uniform([images.shape[0]],
maxval=height_to_pad,
dtype=tf.int32)
width_offsets = tf.random.uniform([images.shape[0]],
maxval=width_to_pad,
dtype=tf.int32)
else:
assert randomness.shape[0] == images.shape[0]
randomness = randomness.numpy() if tf.is_tensor(
randomness) else randomness
height_offsets = randomness % height_to_pad
width_offsets = (randomness // height_to_pad) % width_to_pad
mesh = tf.meshgrid(
tf.range(images.shape[0]),
tf.range(patch_height),
tf.range(patch_width),
indexing="ij",
)
mesh[1] += height_offsets[:, None, None]
mesh[2] += width_offsets[:, None, None]
indices = tf.stack(mesh, axis=-1)
# Paste patches into the images at random locations
return tf.tensor_scatter_nd_update(images, indices, patches)
def _remove_input_items_from_prediction(batch_item_indices, result):
""" Remove input items from predictions, as we don't want to predict them. It is done setting their
probabilities to -1
Args:
batch_item_indices: Batch input item indices
result: Batch predicted probabilities
Returns:
Batch predicted probabilities with input items probs. set to -1
"""
# batch_item_indices is a ragged with input indices for each batch row. Ex [ [0] , [1, 2] ]
batch_item_indices = batch_item_indices.to_tensor(
-1) # -> [ [0,-1] , [1, 2] ]
#print(batch_item_indices, batch_item_indices.shape[0])
# Convert batch_item_indices row indices to (row,column) indices
row_indices = tf.range(0,
tf.shape(batch_item_indices)[0],
dtype=tf.int64) # -> [ 0, 1 ]
row_indices = tf.repeat(
row_indices,
[tf.shape(batch_item_indices)[1]]) # -> [ 0, 0, 1, 1 ]
#print(">>>", batch_item_indices)
batch_item_indices = tf.reshape(batch_item_indices,
shape=[-1]) # -> [ 0, -1, 1, 2 ]
batch_item_indices = tf.stack(
[row_indices, batch_item_indices],
axis=1) # -> [ [0,0] , [0,-1], [1,1], [1,2] ]
# batch_item_indices.to_tensor(-1) added -1's to pad the matrix. Remove these indices
# Needed according to tf.tensor_scatter_nd_update doc. (it will fail in CPU execution, if there are out of bound indices)
# Get indices without -1's:
gather_idxs = tf.where(
batch_item_indices[:, 1] != -1) # -> [[0], [2], [3]]
batch_item_indices = tf.gather_nd(
batch_item_indices, gather_idxs) # -> [ [0,0] , [1,1], [1,2] ]
# To remove input indices, we will set a probability -1 in their indices
updates = tf.repeat(
-1.0,
tf.shape(batch_item_indices)[0]) # -> [ -1, -1, -1 ]
# Assign -1's to the input indices:
return tf.tensor_scatter_nd_update(result, batch_item_indices, updates)
def _update_tp_fp_per_image(self, ground_truth_boxes, predicted_boxes):
with tf.name_scope("update_tp_fp_per_image"):
ground_truth_boxes = tf.boolean_mask(ground_truth_boxes,
tf.logical_not(tf.reduce_all(ground_truth_boxes == 0, 1)))
predicted_boxes = tf.boolean_mask(predicted_boxes,
tf.logical_not(tf.reduce_all(predicted_boxes == 0, 1)))
num_gt_boxes = tf.shape(ground_truth_boxes)[0]
num_pred_boxes = tf.shape(predicted_boxes)[[0]]
matched_gt_boxes = tf.zeros([tf.shape(ground_truth_boxes)[0]], dtype=tf.int32)
self.num_gt_boxes.assign_add(tf.cast(num_gt_boxes, self.num_gt_boxes.dtype))
gt_areas = ((ground_truth_boxes[:, 2] - ground_truth_boxes[:, 0]) *
(ground_truth_boxes[:, 3] - ground_truth_boxes[:, 1]))
pred_areas = ((predicted_boxes[:, 2] - predicted_boxes[:, 0]) *
(predicted_boxes[:, 3] - predicted_boxes[:, 1]))
if tf.greater(tf.shape(predicted_boxes)[0], 0):
for i in tf.range(num_pred_boxes):
box = predicted_boxes[i]
inter_y1 = tf.math.maximum(box[0], ground_truth_boxes[:, 0])
inter_x1 = tf.math.maximum(box[1], ground_truth_boxes[:, 1])
inter_y2 = tf.math.minimum(box[2], ground_truth_boxes[:, 2])
inter_x2 = tf.math.minimum(box[3], ground_truth_boxes[:, 3])
inter_h = tf.math.maximum(inter_y2 - inter_y1, 0.0)
inter_w = tf.math.maximum(inter_x2 - inter_x1, 0.0)
inter_areas = inter_h * inter_w
iou = inter_areas / (gt_areas + pred_areas[i] - inter_areas)
max_iou = tf.reduce_max(iou)
if tf.greater_equal(max_iou, self.iou_threshold):
arg_max_iou = tf.argmax(iou)
if tf.not_equal(matched_gt_boxes[arg_max_iou], 1):
self.true_positives.assign_add(1)
matched_gt_boxes = tf.tensor_scatter_nd_update(matched_gt_boxes,
tf.reshape(arg_max_iou, [1, 1]),
tf.constant([1], dtype=matched_gt_boxes.dtype))
else:
self.false_negatives.assign_add(1)
else:
self.false_negatives.assign_add(1)
def __call__(self, confs, locs, gt_confs, gt_locs):
""" Compute losses for SSD
regression loss: smooth L1
classification loss: cross entropy
Args:
confs: outputs of classification heads (B, num_default, num_classes)
locs: outputs of regression heads (B, num_default, 4)
gt_confs: classification targets (B, num_default)
gt_locs: regression targets (B, num_default, 4)
Returns:
conf_loss: classification loss
loc_loss: regression loss
"""
cross_entropy = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction='none')
# compute classification losses
# without reduction
temp_loss = cross_entropy(gt_confs, confs)
pos = gt_confs > 0
temp_loss = tf.tensor_scatter_nd_update(
temp_loss, tf.where(pos), tf.zeros(tf.shape(tf.where(pos))[0]))
pos_idx, neg_idx = hard_negative_mining(temp_loss, gt_confs,
self.neg_ratio)
# classification loss will consist of positive and negative examples
cross_entropy = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction='sum')
smooth_l1_loss = tf.keras.losses.Huber(reduction='sum')
conf_loss = cross_entropy(
gt_confs[tf.math.logical_or(pos_idx, neg_idx)],
confs[tf.math.logical_or(pos_idx, neg_idx)])
# regression loss only consist of positive examples
loc_loss = smooth_l1_loss(gt_locs[pos_idx], locs[pos_idx])
num_pos = tf.reduce_sum(tf.dtypes.cast(pos_idx, tf.float32))
conf_loss = conf_loss / num_pos
loc_loss = loc_loss / num_pos
return conf_loss, loc_loss
def _drop_nodes_helper(
features: Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]
) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
"""Performs batch-wise operation to drop nodes and corresponding edges."""
atoms, atom_mask, pairs, pair_mask = features
# Select nodes to drop.
idx_drop = self._sample_nodes(atom_mask)
if tf.math.equal(tf.shape(idx_drop)[0], 0):
return features
idx_dropped_nodes.append(tf.squeeze(idx_drop))
# Drop selected nodes.
aug_atom_mask = tf.tensor_scatter_nd_update(
atom_mask, idx_drop, tf.zeros(tf.shape(idx_drop)[0]))
if self.perturb_node_features:
aug_atoms = tf.tensor_scatter_nd_update(
atoms, idx_drop, tf.zeros(
(tf.shape(idx_drop)[0], tf.shape(atoms)[1])))
else:
aug_atoms = atoms
# Remove edges attached to dropped nodes.
# First, remove edges where source node has been dropped (i.e. set rows
# of dropped node indices to 0).
if self.perturb_edge_features:
aug_pairs = tf.tensor_scatter_nd_update(
pairs, idx_drop,
tf.zeros((tf.shape(idx_drop)[0], tf.shape(pairs)[1],
tf.shape(pairs)[-1])))
else:
aug_pairs = pairs
aug_pair_mask = tf.tensor_scatter_nd_update(
pair_mask, idx_drop,
tf.zeros((tf.shape(idx_drop)[0], tf.shape(pair_mask)[-1])))
# Second, remove edges where target node has been dropped (i.e. set
# columns of dropped node indices to 0).
columns = idx_drop
rows = tf.range(tf.shape(pair_mask)[0], dtype=tf.int64)
ii, jj = tf.meshgrid(rows, columns, indexing='ij')
idx_to_update = tf.stack([ii, jj], axis=-1)
updated_values_pair_mask = tf.broadcast_to(0., tf.shape(ii))
aug_pair_mask = tf.tensor_scatter_nd_update(aug_pair_mask, idx_to_update,
updated_values_pair_mask)
if self.perturb_edge_features:
updated_values_pairs = tf.zeros(
(tf.shape(ii)[0], tf.shape(ii)[1], pairs.shape[-1]))
aug_pairs = tf.tensor_scatter_nd_update(aug_pairs, idx_to_update,
updated_values_pairs)
return aug_atoms, aug_atom_mask, aug_pairs, aug_pair_mask
def transform_targets_for_output(y_true, grid_size, anchor_idxs):
# y_true: (N, boxes, (x1, y1, x2, y2, class, best_anchor))
N = tf.shape(y_true)[0]
y_true_out = tf.zeros(
(N, grid_size, grid_size, tf.shape(anchord_idxs)[0], 6)
)
anchor_idxs = tf.cast(anchor_idxs, tf.int32)
indexes = tf.TensorArray(tf.int32, 1, dynamic_size=True)
updates = tf.TensorArray(tf.float32, 1, dynamic_size=True)
idx = 0
for i in tf.range(N):
for j in tf.range(tf.shape(y_true)[1]):
if tf.equal(y_true[i][j][2], 0):
continue
anchor_eq = tf.equal(
anchor_idxs, tf.cast(y_true[i][j][5], tf.int32)
)
if tf.reduce_any(anchor_eq):
box = y_true[i][j][0:4]
box_xy = (y_true[i][j][0:2] + y_true[i][j][2:4])
anchor_idx = tf.cast(tf.where(anchor_eq), tf.int32)
grid_xy = tf.cast(box_xy // (1 / grid_size), tf.int32)
indexes = indexes.write(
idx, [i, grid_xy[1], grid_xy[0], anchor_idx[0][0]]
)
updates = updates.write(
idx, [box[0], box[1], box[2], box[3], 1, y_true[i][j][4]]
)
idx += 1
return tf.tensor_scatter_nd_update(
y_true_out, indexes.stack(), updates.stack()
)
def _compare_tensor_scatter_nd_update_dynamic_indices_with_tf(
test_case,
params_shape,
indices_shape,
updates_shape,
indices_static_shape,
updates_static_shape,
):
params, updates, indices = _random_inputs(params_shape, indices_shape,
updates_shape, False)
i = tf.constant(indices)
x = tf.Variable(params)
y = tf.Variable(updates)
z = tf.tensor_scatter_nd_update(x, i, y)
of_z = _of_tensor_scatter_nd_update_dynamic_indices(
params, indices, updates, indices_static_shape, updates_static_shape)
test_case.assertTrue(np.allclose(z.numpy(), of_z))
def kmeans(self, max_iter, box_num, clusters, k):
dists = tf.zeros((box_num, k))
last = tf.zeros((box_num, ), dtype=tf.int64)
tf.print(tf.shape(clusters))
num_iters = 0
while tf.math.less(num_iters, max_iter):
dists = 1 - self.iou(self._boxes, clusters)
curr = tf.math.argmin(dists, axis=-1)
if tf.math.reduce_all(curr == last):
break
for i in range(k):
hold = tf.math.reduce_mean(self._boxes[curr == i], axis=0)
clusters = tf.tensor_scatter_nd_update(clusters, [[i]], [hold])
last = curr
num_iters += 1
tf.print('k-Means box generation iteration: ', num_iters, end='\r')
return clusters
def group_dist_mat_by_label(i, labels, dist_mat, grouped_dist):
# half of batch
num_triplet_pair_per_sample = tf.math.floordiv(tf.size(labels), 2)
idx = tf.squeeze(tf.where(tf.equal(labels, i)), axis=1)
if tf.shape(idx)[0] == 0:
if i == 0.:
grouped_dist = tf.zeros((1, num_triplet_pair_per_sample))
else:
grouped_dist = tf.concat(
[grouped_dist,
tf.zeros((1, num_triplet_pair_per_sample))],
axis=0)
i = i + 1
return i, labels, dist_mat, grouped_dist
probability = tf.zeros([tf.size(labels)], dtype=tf.float32) # batch size
updates = tf.fill(tf.shape(idx),
tf.truediv(1.0, tf.cast(tf.size(idx), tf.float32)))
probability = tf.tensor_scatter_nd_update(probability,
tf.expand_dims(idx, 1), updates)
probability
triplet_pair_idx = tf.squeeze(tf.random.categorical(
tf.math.log(tf.expand_dims(probability, 0)),
num_triplet_pair_per_sample),
axis=0)
triplet_pair = tf.squeeze(tf.gather(dist_mat,
triplet_pair_idx,
name="triplet_pair_gather"),
axis=1)
if i == 0.:
grouped_dist = tf.expand_dims(triplet_pair, axis=0)
else:
grouped_dist = tf.concat(
[grouped_dist, tf.expand_dims(triplet_pair, axis=0)], axis=0)
i = i + 1
return i, labels, dist_mat, grouped_dist
def _update_context(self, context_in, V, scores, index, L_Q):
B, H, L_V, D = V.shape
if self.mask_flag:
attn_mask = ProbMask(B, H, L_Q, index, scores)
# scores.masked_fill_(attn_mask.mask, -np.inf)
num = 3.4 * math.pow(10, 38)
scores = (scores * attn_mask.mask) + (-((attn_mask.mask * num + num) - num))
attn = tf.keras.activations.softmax(scores, axis=-1) # nn.Softmax(dim=-1)(scores)
batch_indexes = tf.tile(tf.range(V.shape[0])[:, tf.newaxis, tf.newaxis], (1, V.shape[1], index.shape[-1]))
head_indexes = tf.tile(tf.range(V.shape[1])[tf.newaxis, :, tf.newaxis], (V.shape[0], 1, index.shape[-1]))
idx = tf.stack(values=[batch_indexes, head_indexes, index], axis=-1)
context_in = tf.tensor_scatter_nd_update(context_in, idx, tf.matmul(attn, V))
return tf.convert_to_tensor(context_in)
def magneticField(self, gaugeField, scalarField, cpt):
cpt1 = (cpt + 1) % 3
cpt2 = (cpt + 2) % 3
magneticField = tf.math.angle(
tf.linalg.trace(
self.u1Plaquette(gaugeField, scalarField, cpt1, cpt2)))
if (cpt == 0 and self.tHooftLine):
# Correct values along the 't Hooft line
latShape = tf.shape(magneticField)
indices = self.tHooftLineIndices(latShape)
updates = tf.gather_nd(magneticField, indices)
updates = updates - np.pi * tf.sign(updates)
magneticField = tf.tensor_scatter_nd_update(
magneticField, indices, updates)
return 2.0 / self.gaugeCoupling * magneticField
def _dequeue_and_enqueue(self, keys):
# keys: [GN, C]
end_queue_ptr = self.queue_ptr + self.global_batch_size
indices = tf.range(self.queue_ptr, end_queue_ptr,
dtype=tf.int64) # [GN, ]
indices = tf.expand_dims(indices, axis=1) # [GN, 1]
# update to new values
updated_queue = tf.tensor_scatter_nd_update(tensor=self.queue,
indices=indices,
updates=keys)
updated_queue_ptr = end_queue_ptr % self.K
# update queue
self.queue.assign(updated_queue)
# update pointer
self.queue_ptr.assign(updated_queue_ptr)
return
def _merge_sorted(a, b, *args):
"""
Merge sorted arrays efficiently, inspired by https://stackoverflow.com/a/54131815
Parameters
----------
a: tf.Tensor
Sorted tensor for ordering
b: tf.Tensor
Sorted tensor for ordering
args: list of tuple of tf.Tensor
Some data ordered according to a and b that need to be merged whilst keeping the order.
Returns
-------
cs: list of tf.Tensor
Merging of a_x and b_x in the right order.
"""
with tf.name_scope("merge_sorted"):
assert len(a.shape) == len(b.shape) == 1
a_shape, b_shape = tf.shape(a)[0], tf.shape(b)[0]
c_len = tf.shape(a)[0] + tf.shape(b)[0]
if a_shape < b_shape:
a, b = b, a
a_shape, b_shape = tf.shape(a)[0], tf.shape(b)[0]
args = tuple((j, i) for i, j in args)
b_indices = tf.range(b_shape, dtype=tf.int32) + tf.searchsorted(a, b)
a_indices = tf.ones((c_len, ), dtype=tf.bool)
a_indices = tf.tensor_scatter_nd_update(
a_indices, b_indices[:, None], tf.zeros_like(b_indices, tf.bool))
c_range = tf.range(c_len, dtype=tf.int32)
a_mask = tf.boolean_mask(c_range, a_indices)[:, None]
def _inner_merge(u, v):
c = tf.concat([u, v], 0)
c = tf.tensor_scatter_nd_update(c, b_indices[:, None], v)
c = tf.tensor_scatter_nd_update(c, a_mask, u)
return c
return (_inner_merge(a, b), ) + tuple(
_inner_merge(i, j) for i, j in args)
def read_one_mol(idx):
atoms = oe_mol_dicts[int(idx.numpy())]['atomic_symbols']
atoms = tf.expand_dims(tf.convert_to_tensor(
atoms,
tf.string),
1)
atoms = tf.cast(
tf.map_fn(
lambda x: TRANSLATION[x.numpy()[0]],
atoms,
tf.int32),
tf.int64)
atoms = tf.reshape(
atoms,
[-1])
n_atoms = tf.shape(atoms, tf.int64)[0]
bonds = tf.convert_to_tensor(
oe_mol_dicts[int(idx.numpy())]['connectivity'],
dtype=tf.float32)
adjacency_map = tf.zeros(
(n_atoms, n_atoms),
tf.float32)
adjacency_map = tf.tensor_scatter_nd_update(
adjacency_map,
tf.cast(
bonds[:, :2],
tf.int64),
bonds[:, 2])
adjacency_map = tf.i_o.utils.conjugate_average(atoms, adjacency_map)
charges = tf.convert_to_tensor(
oe_mol_dicts[idx]['partial_charges'],
tf.float32)
return atoms, adjacency_map, charges
def test2():
'''测试不变卷积核'''
kernel_value = tf.zeros((3, 3, 2, 4), dtype=tf.float32)
for i in range(4):
kernel_value = tf.tensor_scatter_nd_update(kernel_value,
[[1, 1, i % 2, i]], [1])
# print('kernel_value:', kernel_value)
x = tf.random.uniform((1, 4, 4, 4), dtype=tf.float32)
conv1 = tf.keras.layers.Conv2D(4,
3,
strides=1,
padding='same',
groups=2,
use_bias=False)
out = conv1(x, training=False)
conv1.kernel.assign(kernel_value)
out = conv1(x, training=False)
tf.print('x:', x)
tf.print('out:', out)
def random_proba(shape, clipindex=None, clipproba=None, seed=None):
logits = tf.random.normal(shape=shape,
mean=0.0,
stddev=xsigma(shape),
dtype=tf.float32,
seed=seed)
if clipindex is not None:
clipproba = clip_for_logits(clipproba)
cliplogits = tf.math.log(clipproba) - tf.math.log(1.0 - clipproba)
logits = tf.tensor_scatter_nd_update(logits, [clipindex], [cliplogits])
proba = tf.math.softmax(logits)
return proba
def single_batch_nms(candidate_boxes):
# filter out predictions with score less than score_threshold
candidate_boxes = tf.boolean_mask(
candidate_boxes,
candidate_boxes[..., 4] >= self.score_threshold)
outputs = tf.zeros(
(self.max_detection + 1, tf.shape(candidate_boxes)[-1]))
indices = []
updates = []
count = 0
# keep running this until there's no more candidate box or max_detection is met
while tf.shape(
candidate_boxes)[0] > 0 and count < self.max_detection:
# pick the box with the highest score
best_idx = tf.math.argmax(candidate_boxes[..., 4], axis=0)
best_box = candidate_boxes[best_idx]
# add this best box to the output
indices.append([count])
updates.append(best_box)
count += 1
# remove this box from candidate boxes
candidate_boxes = tf.concat([
candidate_boxes[0:best_idx],
candidate_boxes[best_idx + 1:tf.shape(candidate_boxes)[0]]
],
axis=0)
# calculate IOU between this box and all remaining candidate boxes
iou = broadcast_iou(best_box[0:4], candidate_boxes[..., 0:4])
# remove all candidate boxes with IOU bigger than iou_threshold
candidate_boxes = tf.boolean_mask(candidate_boxes,
iou[0] <= self.iou_threshold)
if count > 0:
# also append num_detection to the result
count_index = [[self.max_detection]]
count_updates = [
tf.fill([tf.shape(candidate_boxes)[-1]], count)
]
indices = tf.concat([indices, count_index], axis=0)
updates = tf.concat([updates, count_updates], axis=0)
outputs = tf.tensor_scatter_nd_update(outputs, indices,
updates)
return outputs
