def loss_infomax(self, x_local, x, edge_index, batch): l_enc = x_local g_enc = x measure = 'JSD' num_graphs = g_enc.shape[0] num_nodes = l_enc.shape[0] pos_mask = torch.zeros((num_nodes, num_graphs)).cuda() neg_mask = torch.ones((num_nodes, num_graphs)).cuda() for nodeidx, graphidx in enumerate(batch): pos_mask[nodeidx][graphidx] = 1. neg_mask[nodeidx][graphidx] = 0. res = torch.mm(l_enc, g_enc.t()) E_pos = get_positive_expectation(res * pos_mask, measure, average=False).sum() E_pos = E_pos / num_nodes E_neg = get_negative_expectation(res * neg_mask, measure, average=False).sum() E_neg = E_neg / (num_nodes * (num_graphs - 1)) loss = E_neg - E_pos return loss
def routine(self, Z_P, measure: str = 'KL'): ''' Args: Z_P: Input unshuffled tensor. ndm_measure: Measure to compare representations with shuffled versions. ''' Z_Q = random_permute(Z_P) E_pos, E_neg, P_samples, Q_samples = self.score( Z_P.detach(), Z_Q.detach(), measure) difference = E_pos - E_neg if measure == 'DV': ndm = E_pos - E_neg else: ndm = get_positive_expectation( P_samples, 'DV') - get_negative_expectation(Q_samples, 'DV') self.add_results(**{measure: difference.detach().item()}) self.add_results(Scores={ 'E_P[D(x)]': P_samples.mean().detach().item(), 'max(D(x))': P_samples.max().detach().item(), 'E_Q[D(x)]': Q_samples.mean().detach().item() }, NDM=ndm.detach().item()) self.add_losses(ndm=-difference)
def local_global_loss_(l_enc, g_enc, edge_index, batch, measure): ''' Args: l: Local feature map. g: Global features. measure: Type of f-divergence. For use with mode `fd` mode: Loss mode. Fenchel-dual `fd`, NCE `nce`, or Donsker-Vadadhan `dv`. Returns: torch.Tensor: Loss. ''' num_graphs = g_enc.shape[0] num_nodes = l_enc.shape[0] pos_mask = torch.zeros((num_nodes, num_graphs)).cuda() neg_mask = torch.ones((num_nodes, num_graphs)).cuda() for nodeidx, graphidx in enumerate(batch): pos_mask[nodeidx][graphidx] = 1. neg_mask[nodeidx][graphidx] = 0. res = torch.mm(l_enc, g_enc.t()) E_pos = get_positive_expectation(res * pos_mask, measure, average=False).sum() E_pos = E_pos / num_nodes E_neg = get_negative_expectation(res * neg_mask, measure, average=False).sum() E_neg = E_neg / (num_nodes * (num_graphs - 1)) return E_neg - E_pos
def fenchel_dual_loss(l, g, measure=None): '''Computes the f-divergence distance between positive and negative joint distributions. Divergences supported are Jensen-Shannon `JSD`, `GAN` (equivalent to JSD), Squared Hellinger `H2`, Chi-squeared `X2`, `KL`, and reverse KL `RKL`. Args: l: Local feature map. g: Global features. measure: f-divergence measure. Returns: torch.Tensor: Loss. ''' N, local_units, n_locs = l.size() l = l.permute(0, 2, 1) l = l.reshape(-1, local_units) u = torch.mm(g, l.t()) u = u.reshape(N, N, -1) mask = torch.eye(N).cuda() n_mask = 1 - mask E_pos = get_positive_expectation(u, measure, average=False).mean(2) E_neg = get_negative_expectation(u, measure, average=False).mean(2) E_pos = (E_pos * mask).sum() / mask.sum() E_neg = (E_neg * n_mask).sum() / n_mask.sum() loss = E_neg - E_pos return loss
def fenchel_dual_loss(l, m, measure=None, neg=False): '''Computes the f-divergence distance between positive and negative joint distributions. Note that vectors should be sent as 1x1. Divergences supported are Jensen-Shannon `JSD`, `GAN` (equivalent to JSD), Squared Hellinger `H2`, Chi-squeared `X2`, `KL`, and reverse KL `RKL`. Args: l: Local feature map. m: Multiple globals feature map. measure: f-divergence measure. Returns: torch.Tensor: Loss. ''' N, units, n_locals = l.size() n_multis = m.size(2) # First we make the input tensors the right shape. l = l.view(N, units, n_locals) l = l.permute(0, 2, 1) l = l.reshape(-1, units) m = m.view(N, units, n_multis) m = m.permute(0, 2, 1) m = m.reshape(-1, units) # Outer product, we want a N x N x n_local x n_multi tensor. u = torch.mm(m, l.t()) u = u.reshape(N, n_multis, N, n_locals).permute(0, 2, 3, 1) # Since we have a big tensor with both positive and negative samples, we need to mask. mask = torch.eye(N).to(l.device) n_mask = 1 - mask # Compute the positive and negative score. Average the spatial locations. E_pos = get_positive_expectation(u, measure, average=False).mean(2).mean(2) E_neg = get_negative_expectation(u, measure, average=False).mean(2).mean(2) # Mask positive and negative terms for positive and negative parts of loss E_pos = (E_pos * mask).sum() / mask.sum() E_neg = (E_neg * n_mask).sum() / n_mask.sum() loss = E_neg - E_pos if neg: return - loss else: return loss
def score(self, X_P, X_Q, measure): '''Score real and fake. Args: X_P: Real tensor (unshuffled). X_Q: Fake tensor (shuffled). measure: Comparison measure. Returns: tuple of torch.Tensor: (real expectation, fake expectation, real samples, fake samples) ''' ndm = self.nets.ndm P_samples = ndm(X_P) Q_samples = ndm(X_Q) E_pos = get_positive_expectation(P_samples, measure) E_neg = get_negative_expectation(Q_samples, measure) return E_pos, E_neg, P_samples, Q_samples