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