The "LocalModel()" used in the model.py file is inconsistent with the "LocalModel()" defined in your local_module.py.
Mang30 opened this issue · comments
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
