def build(self, inputs, targets):
"""
Args:
inputs (torch.Tensor): feature matrix with shape (batch_size, feat_dim).
targets (torch.LongTensor): ground truth labels with shape (num_classes).
"""
n = inputs.shape[0]
dist = math.reduce_sum(math.pow(
inputs, flow.constant_like(inputs, 2, dtype=flow.float32)),
axis=1)
shape_tensor = flow.constant(value=0.0,
dtype=flow.float32,
shape=(n, n))
dist = flow.broadcast_like(dist, like=shape_tensor, broadcast_axes=[1])
dist = math.add(
dist, flow.transpose(dist, perm=(1, 0),
batch_axis_non_change=True))
temp1 = math.multiply(
-2,
flow.matmul(
inputs,
flow.transpose(inputs, perm=(1, 0),
batch_axis_non_change=True)))
dist = math.add(dist, temp1)
dist = math.sqrt(flow.clamp(dist, min_value=1e-12))
mask = math.equal(
flow.broadcast_like(targets, like=shape_tensor,
broadcast_axes=[1]),
flow.transpose(flow.broadcast_like(targets,
like=shape_tensor,
broadcast_axes=[1]),
perm=(1, 0),
batch_axis_non_change=True))
mask_rev = math.not_equal(
flow.broadcast_like(targets, like=shape_tensor,
broadcast_axes=[1]),
flow.transpose(flow.broadcast_like(targets,
like=shape_tensor,
broadcast_axes=[1]),
perm=(1, 0),
batch_axis_non_change=True))
dist_ap, dist_an = [], []
for i in range(n):
temp_dist = flow.slice_v2(dist, [(i, i + 1, 1)])
temp_mask = flow.slice_v2(mask, [(i, i + 1, 1)])
temp_mask_rev = flow.slice_v2(mask_rev, [(i, i + 1, 1)])
dist_ap.append(
math.reduce_max(
flow.gather_nd(temp_dist, flow.where(temp_mask))))
dist_an.append(
math.reduce_min(
flow.gather_nd(temp_dist, flow.where(temp_mask_rev))))
dist_ap = flow.concat(dist_ap, 0)
dist_an = flow.concat(dist_an, 0)
y = flow.ones_like(dist_an)
# return dist_an, dist_ap, y
return self._MarginRankingLoss(dist_an, dist_ap, y)
def gather_nd_fn(
params_def: oft.ListNumpy.Placeholder(params.shape, dtype=flow.float),
indices_def: oft.ListNumpy.Placeholder(indices_static_shape,
dtype=flow.int32),
):
with flow.scope.placement(device_type, "0:0"):
return flow.gather_nd(params_def, indices_def)
def _test_gather_nd_t(test_case, device):
input = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
indices = np.array([[0, 2], [2, 1]])
np_out = np.array([3, 8])
output = flow.gather_nd(
flow.tensor(input, dtype=flow.float, device=flow.device(device)),
flow.tensor(indices, dtype=flow.int, device=flow.device(device)),
)
test_case.assertTrue(np.array_equal(output.numpy(), np_out))
def do_gather_nd(x, index):
x_var = flow.get_variable(
"params", shape=(1,), dtype=x_dtype, initializer=flow.zeros_initializer(),
)
x = x + flow.cast_to_current_logical_view(x_var)
y = flow.gather_nd(x, index)
if need_grad:
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(y)
if callable(comp_diff_fn):
flow.watch_diff(x, comp_diff_fn)
return y
def gather_nd():
x = flow.get_variable(
name="x",
shape=(2, 3, 4),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
y = flow.get_variable(
name="y",
shape=(2, 3),
dtype=flow.int64,
initializer=flow.random_uniform_initializer(0, 1, flow.int64),
)
return flow.gather_nd(x, y)
def do_gather_nd(x_blob, i_blob):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"params",
shape=params.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x = flow.cast_to_current_logical_view(x)
x = x + x_blob
y = flow.gather_nd(x, i_blob)
flow.losses.add_loss(y)
flow.watch_diff(x, compare_fn)
return y
def _test_gather_nd_backward_t(test_case, device):
input = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
indices = np.array([[0, 2], [2, 1]])
np_out = np.array([3, 8])
np_grad = np.array([[0, 0, 1], [0, 0, 0], [0, 1, 0]])
of_input = flow.tensor(input,
requires_grad=True,
dtype=flow.float,
device=flow.device(device))
output = flow.gather_nd(
of_input,
flow.tensor(indices, dtype=flow.int, device=flow.device(device)))
out_sum = output.sum()
out_sum.backward()
test_case.assertTrue(np.array_equal(output.numpy(), np_out))
test_case.assertTrue(np.array_equal(of_input.grad.numpy(), np_grad))
def do_gather_nd(x_blob, i_blob):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"params",
shape=params.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x = flow.cast_to_current_logical_view(x)
x = x + x_blob
y = flow.gather_nd(x, i_blob)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(y)
flow.watch_diff(x, compare_fn)
return y
def gather_nd_fn(
params_def: oft.ListNumpy.Placeholder(static_params_shape, dtype=flow.float),
indices_def: oft.ListNumpy.Placeholder(indices.shape, dtype=flow.int32),
):
with flow.scope.placement("gpu", "0:0"):
one_var = flow.get_variable(
"one",
shape=(1,),
dtype=flow.float32,
initializer=flow.constant_initializer(1),
)
one_var = flow.cast_to_current_logical_view(one_var)
params_var = params_def * one_var
y = flow.gather_nd(params_var, indices_def)
flow.losses.add_loss(y)
flow.watch_diff(params_var, compare_fn)
return y
def gather_nd_fn(
params_def: oft.ListNumpy.Placeholder(static_params_shape, dtype=flow.float),
indices_def: oft.ListNumpy.Placeholder(indices.shape, dtype=flow.int32),
):
with flow.scope.placement("gpu", "0:0"):
one_var = flow.get_variable(
"one",
shape=(1,),
dtype=flow.float32,
initializer=flow.constant_initializer(1),
)
one_var = flow.cast_to_current_logical_view(one_var)
params_var = params_def * one_var
y = flow.gather_nd(params_var, indices_def)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(y)
flow.watch_diff(params_var, compare_fn)
return y
def __call__(self, x, padding=None):
# Retrieve dynamically known shapes
batch_size = x.shape[0]
length = x.shape[1]
if padding is not None:
with flow.scope.namespace("remove_padding"):
# Flatten padding to [batch_size*length]
pad_mask = flow.reshape(padding, [-1])
nonpad_ids = flow.cast(flow.where(pad_mask < 1e-9), dtype=flow.int32)
# nonpad_ids = tf.to_int32(tf.where(pad_mask < 1e-9))
# Reshape x to [batch_size*length, hidden_size] to remove padding
x = flow.reshape(x, [-1, self.hidden_size])
x = flow.gather_nd(x, indices=nonpad_ids)
# Reshape x from 2 dimensions to 3 dimensions.
# TODO:Maybe has a batch axis error in there
x = flow.expand_dims(x, axis=0)
output = self._build_dense(x, self.filter_size, name="filter_layer")
if self.train:
# In TensorFlow the param means `keep_prob` and use `1-dropout`,
# but our dropout means drop rate so i just use dropout !
output = flow.nn.dropout(output, self.relu_dropout)
if padding is not None:
with flow.scope.namespace("re_add_padding"):
output = flow.squeeze(output, axis=[0, ])
output = flow.scatter_nd(
indices=nonpad_ids,
updates=output,
shape=[batch_size * length, self.hidden_size]
)
output = flow.reshape(output, [batch_size, length, self.hidden_size])
return output
def forward(self, inputs, targets):
n = inputs.shape[0]
# Compute pairwise distance, replace by the official when merged
tempname = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S.%f')
shape_tensor = flow.constant(value=0.0,
dtype=flow.float32,
shape=(n, n))
if self.distance == 'euclidean':
blob_2 = flow.get_variable(
"blob_2_" + tempname,
shape=inputs.shape,
initializer=flow.constant_initializer(2),
dtype=inputs.dtype)
dist = flow.math.pow(inputs, blob_2)
dist = flow.math.reduce_sum(dist, axis=1, keepdims=True)
dist = flow.broadcast_like(dist, shape_tensor)
tempdist = flow.transpose(dist)
dist = dist + tempdist
inputs_t = flow.transpose(inputs)
dist = addmm(dist, inputs, inputs_t, beta=1, alpha=-2)
dist = flow.clamp(dist, min_value=1e-12)
dist = flow.math.sqrt(dist)
elif self.distance == 'cosine':
#fnorm=flow.math.l2_normalize(inputs, axis=1)
fnorm = flow.math.reduce_mean(flow.math.divide(
inputs, flow.math.l2_normalize(inputs, axis=1)),
axis=1,
keepdims=True)
expand_fnorm = flow.broadcast_like(fnorm,
like=inputs,
broadcast_axes=[1])
l2norm = flow.math.divide(inputs, expand_fnorm)
l2norm_t = flow.transpose(l2norm, perm=(1, 0))
dist = flow.math.negative(flow.matmul(l2norm, l2norm_t))
# For each anchor, find the hardest positive and negative
mask = math.equal(
flow.broadcast_like(targets, like=shape_tensor,
broadcast_axes=[1]),
flow.transpose(flow.broadcast_like(targets,
like=shape_tensor,
broadcast_axes=[1]),
perm=(1, 0),
batch_axis_non_change=True))
mask_rev = math.not_equal(
flow.broadcast_like(targets, like=shape_tensor,
broadcast_axes=[1]),
flow.transpose(flow.broadcast_like(targets,
like=shape_tensor,
broadcast_axes=[1]),
perm=(1, 0),
batch_axis_non_change=True))
dist_ap, dist_an = [], []
for i in range(n):
temp_dist = flow.slice_v2(dist, [(i, i + 1, 1)])
temp_mask = flow.slice_v2(mask, [(i, i + 1, 1)])
temp_mask_rev = flow.slice_v2(mask_rev, [(i, i + 1, 1)])
temp_dist_ap = flow.expand_dims(
math.reduce_max(
flow.gather_nd(temp_dist, flow.where(temp_mask))), 0)
temp_dist_an = flow.expand_dims(
math.reduce_min(
flow.gather_nd(temp_dist, flow.where(temp_mask_rev))), 0)
dist_ap.append(temp_dist_ap)
dist_an.append(temp_dist_an)
dist_ap = flow.concat(dist_ap, 0)
dist_an = flow.concat(dist_an, 0)
y = flow.ones_like(dist_an)
return self._MarginRankingLoss(dist_an, dist_ap, y)
