panda0406 / zero-shot-knowledge-graph-relational-learning Goto Github PK

    self.gcn_w = nn.Linear(self.embed_dim, int(self.embed_dim/2))
    self.gcn_b = nn.Parameter(torch.FloatTensor(self.embed_dim))

    self.fc1 = nn.Linear(self.embed_dim, int(self.embed_dim/2))
    self.fc2 = nn.Linear(self.embed_dim, int(self.embed_dim/2))

    self.dropout = nn.Dropout(0.2)
    self.dropout_e = nn.Dropout(0.2)

    self.symbol_emb.weight.data.copy_(torch.from_numpy(embed))

    self.symbol_emb.weight.requires_grad = False

    d_model = self.embed_dim * 2
    self.support_encoder = SupportEncoder(d_model, 2*d_model, dropout=0.2)
    #self.query_encoder = QueryEncoder(d_model, process_steps)

def neighbor_encoder(self, connections, num_neighbors):
    '''
    connections: (batch, 200, 2)
    num_neighbors: (batch,)
    '''
    num_neighbors = num_neighbors.unsqueeze(1)
    entities = connections[:,:,1].squeeze(-1)
    ent_embeds = self.dropout(self.symbol_emb(entities)) # (batch, 50, embed_dim)
    concat_embeds = ent_embeds

    out = self.gcn_w(concat_embeds)
    out = torch.sum(out, dim=1) # (batch, embed_dim)
    out = out / num_neighbors
    return out.tanh()

def entity_encoder(self, entity1, entity2):
    entity1 = self.dropout_e(entity1)
    entity2 = self.dropout_e(entity2)
    entity1 = self.fc1(entity1)
    entity2 = self.fc2(entity2)
    entity = torch.cat((entity1, entity2), dim=-1)
    return entity.tanh() # (batch, embed_dim)


def forward(self, query, support, query_meta=None, support_meta=None):
    '''
    query: (batch_size, 2)
    support: (few, 2)
    return: (batch_size, )
    '''
    query_left_connections, query_left_degrees, query_right_connections, query_right_degrees = query_meta
    support_left_connections, support_left_degrees, support_right_connections, support_right_degrees = support_meta

    query_e1 = self.symbol_emb(query[:,0]) # (batch, embed_dim)
    query_e2 = self.symbol_emb(query[:,1]) # (batch, embed_dim)
    query_e = self.entity_encoder(query_e1, query_e2)

    support_e1 = self.symbol_emb(support[:,0]) # (batch, embed_dim)
    support_e2 = self.symbol_emb(support[:,1]) # (batch, embed_dim)
    support_e = self.entity_encoder(support_e1, support_e2)

    query_left = self.neighbor_encoder(query_left_connections, query_left_degrees)
    query_right = self.neighbor_encoder(query_right_connections, query_right_degrees)

    support_left = self.neighbor_encoder(support_left_connections, support_left_degrees)
    support_right = self.neighbor_encoder(support_right_connections, support_right_degrees)
    
    query_neighbor = torch.cat((query_left, query_e,  query_right), dim=-1) # tanh
    support_neighbor = torch.cat((support_left, support_e, support_right), dim=-1) # tanh

    support = support_neighbor
    query = query_neighbor

    support_g = self.support_encoder(support) # 1 * 100
    query_g = self.support_encoder(query)

    support_g = torch.mean(support_g, dim=0, keepdim=True)

    # cosine similarity
    matching_scores = torch.matmul(query_g, support_g.t()).squeeze()

    return query_g, matching_scores