snap-research / largegt Goto Github PK

    self.local_module = LocalModule(
        seq_len=self.sample_node_len * 3,
        input_dim=in_channels,
        n_layers=1,
        num_heads=heads,
        hidden_dim=out_channels,
    )

def local_forward(self, seq, pos_enc):
    return self.local_module(seq, pos_enc)

def forward(self, batched_data):
    tensor = self.att_embeddings_nope(batched_data)

    # transformer encoder
    for enc_layer in self.layers:
        tensor = enc_layer(tensor)

    output = self.final_ln(tensor)

    _target = output[:, 0, :].unsqueeze(1).repeat(1, self.seq_len - 1, 1)
    split_tensor = torch.split(output, [1, self.seq_len - 1], dim=1)

    node_tensor = split_tensor[0]
    _neighbor_tensor = split_tensor[1]

    if self.node_only_readout:
        # only slicing the indices that belong to nodes and not the 1-hop and 2-hop feats
        indices = torch.arange(3, self.seq_len, 3)
        neighbor_tensor = _neighbor_tensor[:, indices]
        target = _target[:, indices]
    else:
        target = _target
        neighbor_tensor = _neighbor_tensor

    layer_atten = self.attn_layer(torch.cat((target, neighbor_tensor), dim=2))
    layer_atten = F.softmax(layer_atten, dim=1)

    neighbor_tensor = neighbor_tensor * layer_atten
    neighbor_tensor = torch.sum(neighbor_tensor, dim=1, keepdim=True)

    output = (node_tensor + neighbor_tensor).squeeze()

    return output