The largegt from snap-research

largegt's People

Contributors

Stargazers

Watchers

largegt's Issues

The "LocalModel()" used in the model.py file is inconsistent with the "LocalModel()" defined in your local_module.py.

Hi, great work you're doing! However, while trying to reproduce the results, I encountered the error "TypeError: LocalModule.forward() takes 2 positional arguments but 3 were given." I guess this is because the "LocalModel()" used in the model.py file is not consistent with the "LocalModel()" defined in your local_module.py.

#####################################################
########## the "LocalModel()" used in the model.py ###########

    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)

#####################################################

the "LocalModel()" defined in your local_module.py.

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