def merge_samples(i_example_to_i_samples, i_example_to_i_samples_2, gauss_samples,
gauss_samples_2):
inds_a = np.sort(np.concatenate(i_example_to_i_samples))
th_inds_a = np_to_var(inds_a, dtype=np.int64)
th_inds_a, _ = ensure_on_same_device(
th_inds_a, gauss_samples)
samples_a = gauss_samples[th_inds_a]
inds_b = np.sort(np.concatenate(i_example_to_i_samples_2))
if len(inds_b) > 0:
th_inds_b = np_to_var(inds_b, dtype=np.int64)
th_inds_b, _ = ensure_on_same_device(
th_inds_b, gauss_samples)
samples_b = gauss_samples_2[th_inds_b]
a_dict = dict([(val, i) for i,val in enumerate(inds_a)])
b_dict = dict([(val, i + len(a_dict)) for i,val in enumerate(inds_b)])
# merge samples
i_example_to_i_samples_merged = []
for i_example in range(len(i_example_to_i_samples)):
a_examples = [a_dict[i] for i in i_example_to_i_samples[i_example]]
b_examples = [b_dict[i] for i in i_example_to_i_samples_2[i_example]]
i_example_to_i_samples_merged.append(a_examples + b_examples)
if len(inds_b) > 0:
all_samples = th.cat((samples_a, samples_b), dim=0)
else:
all_samples = samples_a
return all_samples, i_example_to_i_samples_merged
def transport_mat_from_diffs(diffs):
transport_mat = ot.emd([], [], var_to_np(diffs))
# sometimes weird low values, try to prevent them
transport_mat = transport_mat * (transport_mat > (1.0 / (diffs.numel())))
transport_mat = np_to_var(transport_mat, dtype=np.float32)
diffs, transport_mat = ensure_on_same_device(diffs, transport_mat)
return transport_mat
def ot_euclidean_transport_mat(samples_a, samples_b):
diffs = samples_a.unsqueeze(1) - samples_b.unsqueeze(0)
diffs = th.sqrt(th.clamp(th.sum(diffs * diffs, dim=2), min=1e-6))
transport_mat = ot.emd([], [], var_to_np(diffs))
# sometimes weird low values, try to prevent them
transport_mat = transport_mat * (transport_mat > (1.0 / (diffs.numel())))
transport_mat = np_to_var(transport_mat, dtype=np.float32)
diffs, transport_mat = ensure_on_same_device(diffs, transport_mat)
return transport_mat
def ot_emd_loss_for_samples(samples_a, samples_b):
diffs = samples_a.unsqueeze(1) - samples_b.unsqueeze(0)
diffs = th.sum(diffs * diffs, dim=2)
transport_mat = ot.emd([], [], var_to_np(diffs))
# sometimes weird low values, try to prevent them
transport_mat = transport_mat * (transport_mat > (1.0 / (diffs.numel())))
transport_mat = np_to_var(transport_mat, dtype=np.float32)
diffs, transport_mat = ensure_on_same_device(diffs, transport_mat)
eps = 1e-6
loss = th.sqrt(th.sum(transport_mat * diffs) + eps)
return loss
def ot_emd_loss(outs, mean, std):
gauss_samples = get_gauss_samples(len(outs), mean, std)
diffs = outs.unsqueeze(1) - gauss_samples.unsqueeze(0)
del gauss_samples
diffs = th.sum(diffs * diffs, dim=2)
transport_mat = ot.emd([], [], var_to_np(diffs))
# sometimes weird low values, try to prevent them
transport_mat = transport_mat * (transport_mat > (1.0 / (diffs.numel())))
transport_mat = np_to_var(transport_mat, dtype=np.float32)
diffs, transport_mat = ensure_on_same_device(diffs, transport_mat)
eps = 1e-6
loss = th.sqrt(th.sum(transport_mat * diffs) + eps)
return loss
def unbalanced_transport_mat_squared_diff(samples_a, samples_b, cover_fraction,
return_diffs=False):
diffs = samples_a.unsqueeze(1) - samples_b.unsqueeze(0)
diffs = th.sum(diffs * diffs, dim=2)
# add dummy point with distance 0 to everything
dummy = th.zeros_like(diffs[0:1,:])
diffs = th.cat((diffs, dummy), dim=0)
a = np.ones(len(samples_a)) / len(samples_a) * cover_fraction
a = np.concatenate((a, [1 - cover_fraction]))
transport_mat = ot.emd(a, [], var_to_np(diffs))
transport_mat = np_to_var(transport_mat, dtype=np.float32)
transport_mat, diffs = ensure_on_same_device(transport_mat, diffs)
if return_diffs:
return transport_mat, diffs
else:
return transport_mat
def get_perturbations(l_out, norm, perturb_dist):
dims = int(np.prod(l_out.size()[1:]))
mean = th.zeros(dims)
std = th.ones(dims)
mean = th.autograd.Variable(mean)
std = th.autograd.Variable(std)
_, mean, std = ensure_on_same_device(l_out, mean, std)
if perturb_dist == 'uniform':
perturbations = get_uniform_samples(l_out.size()[0], mean, std)
else:
assert perturb_dist == 'gaussian'
perturbations = get_gauss_samples(l_out.size()[0], mean, std)
perturbations = norm * (perturbations / th.sqrt(th.sum(perturbations ** 2, dim=1, keepdim=True)))
perturbations = perturbations.view(l_out.size())
return perturbations
def get_batch(
inputs,
targets,
rng,
batch_size,
with_replacement,
i_class='all',
):
if i_class == 'all':
indices = list(range(len(inputs)))
else:
indices = np.flatnonzero(var_to_np(targets[:, i_class]) == 1)
batch_inds = rng.choice(indices, size=batch_size, replace=with_replacement)
th_inds = np_to_var(batch_inds, dtype=np.int64)
th_inds, _ = ensure_on_same_device(th_inds, inputs)
batch_X = inputs[th_inds]
batch_y = targets[th_inds]
return th_inds, batch_X, batch_y
def ot_euclidean_loss(outs, mean, std, normalize_by_global_emp_std=False):
gauss_samples = get_gauss_samples(len(outs), mean, std)
diffs = outs.unsqueeze(1) - gauss_samples.unsqueeze(0)
del gauss_samples
if normalize_by_global_emp_std:
global_emp_std = th.mean(th.std(outs, dim=0))
diffs = diffs / global_emp_std
diffs = th.sqrt(th.clamp(th.sum(diffs * diffs, dim=2), min=1e-6))
transport_mat = ot.emd([], [], var_to_np(diffs))
# sometimes weird low values, try to prevent them
transport_mat = transport_mat * (transport_mat > (1.0 / (diffs.numel())))
transport_mat = np_to_var(transport_mat, dtype=np.float32)
diffs, transport_mat = ensure_on_same_device(diffs, transport_mat)
loss = th.sum(transport_mat * diffs)
return loss
def get_features_shuffled(x):
randperms = th.autograd.Variable(
th.stack([th.randperm(len(x)) for _ in range(x.shape[1])], dim=1))
randperms, x = ensure_on_same_device(randperms, x)
shuffled_ins = x.gather(index=randperms, dim=0)
return shuffled_ins