Method get_loss

code/model.py:121–149

(self, x, x_aug, batch_size)

    
    
    def get_loss(self, x, x_aug, batch_size):

        T = 0.2
        n_samples, _ = x.size()
        x_abs = x.norm(dim=1)
        x_aug_abs = x_aug.norm(dim=1)
        if n_samples <= batch_size:
            batch_size = n_samples
            sim_matrix = torch.einsum('ik,jk->ij', x, x_aug) / torch.einsum('i,j->ij', x_abs, x_aug_abs)
            sim_matrix = torch.exp(sim_matrix / T)
            pos_sim = sim_matrix[range(batch_size), range(batch_size)]
            loss = pos_sim / (sim_matrix.sum(dim=1) - pos_sim)
            loss = - torch.log(loss).mean()
        else:
            n_loop = n_samples // batch_size + 1
            losses = []
            for i in range(n_loop):
                start = i*batch_size
                end = (i + 1)*batch_size if i != n_loop - 1 else n_samples
                n_sim = batch_size if i != n_loop - 1 else end - start
                sim_matrix = torch.einsum('ik,jk->ij', x[start:end], x_aug) / torch.einsum('i,j->ij', x_abs[start:end], x_aug_abs)
                sim_matrix = torch.exp(sim_matrix / T)
                pos_sim = sim_matrix[range(n_sim), range(n_sim)]
                loss = pos_sim / (sim_matrix.sum(dim=1) - pos_sim)
                losses.append(-torch.log(loss))
            
            loss = torch.concat(losses).mean()

        return loss
    

class GraphInfoMax(nn.Module):