Running tgn.py on a small dataset

[SalvishGoomanee]

I am trying to run the tgn.py file on a small dataset shown below in order to understand how things work with this particular temporal GNN.

Here is the data:

data = TemporalData(
    src=Tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
    dst=Tensor([22, 33, 44, 55, 66, 77, 88, 99, 100, 111]),
    t=Tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90]),
    msg=Tensor([1, 1, 0, 0, 0, 1, 2, 3, 2, 1])
)
data = data.to(device)

I have not change anything in the code except that I exclude the train/test/val_loader part as the above dataset is small.

I end up with the following error:

IndexError: tensors used as indices must be long, byte or bool tensors at the line n_id, edge_index, e_id = neighbour_loader(n_id) in the training loop.

I am not sure how to correct this.

For the sake of clarity here is the whole script:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

"""Define the temporal data"""
data = TemporalData(
    src=Tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
    dst=Tensor([22, 33, 44, 55, 66, 77, 88, 99, 100, 111]),
    t=Tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90]),
    msg=Tensor([1, 1, 0, 0, 0, 1, 2, 3, 2, 1])
)
data = data.to(device)

# Add property to data as such.
data.surface_tensions = Tensor([1, 10, 20, 30, 15, 20, 45, 10, 23, 25])
print(data)

# Ensure dst nodes are sampled as negatives.
min_dst_idx, max_dst_idx = int(data.dst.min()), int(data.dst.max())
train_data, val_data, test_data = data.train_val_test_split(val_ratio=0.15, test_ratio=0.15)

train_loader = TemporalDataLoader(train_data, batch_size=3)
test_loader = TemporalDataLoader(test_data, batch_size=3)
val_loader = TemporalDataLoader(val_data, batch_size=3)
neighbour_loader = LastNeighborLoader(data.num_nodes, size=2, device=device)

# Follows the TGAT model for embedding
class GraphAttentionEmbedding(torch.nn.Module):
    def __init__(self, in_channels, out_channels, msg_dim, time_enc):
        super().__init__()
        self.time_enc = time_enc
        edge_dim = msg_dim + time_enc.out_channels
        self.conv = TransformerConv(in_channels, out_channels // 2, heads=2,
                                    dropout=0.1, edge_dim=edge_dim)

    def forward(self, x, last_update, edge_index, t, msg):
        rel_t = last_update[edge_index[0]] - t
        rel_t_enc = self.time_enc(rel_t.to(x.dtype))
        edge_attr = torch.cat([rel_t_enc, msg], dim=-1)
        return self.conv(x, edge_index, edge_attr)


class LinkPredictor(torch.nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.lin_src = Linear(in_channels, in_channels)
        self.lin_dst = Linear(in_channels, in_channels)
        self.lin_final = Linear(in_channels, 1)

    def forward(self, z_src, z_dst):
        h = self.lin_src(z_src) + self.lin_dst(z_dst)
        h = h.relu()
        return self.lin_final(h)

memory_dim = time_dim = embedding_dim = 100

memory = TGNMemory(
    data.num_nodes,
    data.msg.size(-1),
    memory_dim,
    time_dim,
    message_module=IdentityMessage(data.msg.size(-1), memory_dim, time_dim),
    aggregator_module=LastAggregator(),
).to(device)

gnn = GraphAttentionEmbedding(
    in_channels=memory_dim,
    out_channels=embedding_dim,
    msg_dim=data.msg.size(-1),
    time_enc=memory.time_enc,
).to(device)

link_pred = LinkPredictor(in_channels=embedding_dim).to(device)

optimizer = torch.optim.Adam(set(memory.parameters()) | set(gnn.parameters()) | set(link_pred.parameters()), lr=0.0001)
criterion = torch.nn.BCEWithLogitsLoss()

# Helper vector to map global node indices to local ones.
assoc = torch.empty(data.num_nodes, dtype=torch.long, device=device)

def train():
    memory.train()
    gnn.train()
    link_pred.train()

    memory.reset_state()  # Start with a fresh memory.
    neighbour_loader.reset_state()  # Start with an empty graph.

    total_loss = 0
    for batch in train_loader:
        batch = batch.to(device)
        optimizer.zero_grad()

        src, pos_dst, t, msg = batch.src, batch.dst, batch.t, batch.msg

        # Sample negative destination nodes.
        neg_dst = torch.randint(min_dst_idx, max_dst_idx + 1, (src.size(0),),
                                dtype=torch.long, device=device)

        n_id = torch.cat([src, pos_dst, neg_dst]).unique()
        n_id, edge_index, e_id = neighbour_loader(n_id)
        print(n_id)
        assoc[n_id] = torch.arange(n_id.size(0), device=device)

        # Get updated memory of all nodes involved in the computation.
        z, last_update = memory(n_id)
        z = gnn(z, last_update, edge_index, data.t[e_id].to(device),
                data.msg[e_id].to(device))

        pos_out = link_pred(z[assoc[src]], z[assoc[pos_dst]])
        neg_out = link_pred(z[assoc[src]], z[assoc[neg_dst]])

        loss = criterion(pos_out, torch.ones_like(pos_out))
        loss += criterion(neg_out, torch.zeros_like(neg_out))

        # Update memory and neighbor loader with ground-truth state.
        memory.update_state(src, pos_dst, t, msg)
        neighbour_loader.insert(src, pos_dst)

        loss.backward()
        optimizer.step()
        memory.detach()
        total_loss += float(loss) * batch.num_events

    return total_loss / train_data.num_events


@torch.no_grad()
def test(loader):
    memory.eval()
    gnn.eval()
    link_pred.eval()

    torch.manual_seed(12345)  # Ensure deterministic sampling across epochs.

    aps, aucs = [], []
    for batch in loader:
        batch = batch.to(device)
        src, pos_dst, t, msg = batch.src, batch.dst, batch.t, batch.msg

        neg_dst = torch.randint(min_dst_idx, max_dst_idx + 1, (src.size(0),),
                                dtype=torch.long, device=device)

        n_id = torch.cat([src, pos_dst, neg_dst]).unique()
        n_id, edge_index, e_id = neighbour_loader(n_id)
        assoc[n_id] = torch.arange(n_id.size(0), device=device)

        z, last_update = memory(n_id)
        z = gnn(z, last_update, edge_index, data.t[e_id].to(device),
                data.msg[e_id].to(device))

        pos_out = link_pred(z[assoc[src]], z[assoc[pos_dst]])
        neg_out = link_pred(z[assoc[src]], z[assoc[neg_dst]])

        y_pred = torch.cat([pos_out, neg_out], dim=0).sigmoid().cpu()
        y_true = torch.cat(
            [torch.ones(pos_out.size(0)),
             torch.zeros(neg_out.size(0))], dim=0)

        aps.append(average_precision_score(y_true, y_pred))
        aucs.append(roc_auc_score(y_true, y_pred))

        memory.update_state(src, pos_dst, t, msg)
        neighbour_loader.insert(src, pos_dst)

    return float(torch.tensor(aps).mean()), float(torch.tensor(aucs).mean())


def run():
    for epoch in range(1, 51):
        loss = train()
        # print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}')
        val_ap, val_auc = test(val_loader)
        test_ap, test_auc = test(test_loader)
        # print(f'Val AP: {val_ap:.4f}, Val AUC: {val_auc:.4f}')
        # print(f'Test AP: {test_ap:.4f}, Test AUC: {test_auc:.4f}')
        print(f'Epoch: {epoch} | Loss: {loss} | Val_ap: {val_ap} | Val_auc: {val_auc} | Test_ap: {test_ap} | Test_auc: {test_auc}')


if __name__ == "__main__":
    run()
    # pass

Thanks in advance for the help!

[rusty1s]

You need to define src, dst and t as LongTensor:

data = TemporalData(
    src=torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
    dst=torch.tensor([22, 33, 44, 55, 66, 77, 88, 99, 100, 111]),
    t=torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90]),
    msg=Tensor([1, 1, 0, 0, 0, 1, 2, 3, 2, 1])
)
data = data.to(device)

[SalvishGoomanee]

This is the whole thing with your last suggestion on making the msg 2D. Here is the same error:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (3x4 and 5x5)

It is pointing to the lines:
h = self.lin_src(z_src) + self.lin_dst(z_dst) in the forward function in LinkPredictor and when first being read in pos_out = link_pred(z[assoc[src]], z[assoc[pos_dst]]) in the for loop of the training function.

[rusty1s]

That's because your embedding_dim is odd, meaning that out_channels // 2 results in a rounding error. Changing embedding_dim=6 resolves this.

[SalvishGoomanee]

Thanks! It works in the training loop. In the test loop, I have the error at this line:
memory.update_state(src, pos_dst, t, msg)

Error: RuntimeError: mat1 and mat2 must have the same dtype at line 182.

Is it because the types in the beginning were changed to LongTensors?

[rusty1s]

You need to use msg=torch.Tensor([1, 1, 0, 0, 0, 1, 2, 3, 2, 1]).view(-1, 1), that works for me.

[SalvishGoomanee]

Yes this is finally working now! Thanks a lot Matthias!

[SalvishGoomanee]

This is in relation to the thread above and relates to training the model in tgn.py with multiple (dynamical) node and edge features. I provide a very small example of such data here which works with the above code in tgn.py updated as per in the discussion that follows.

The questions are in the bullet points.

I have a dataset stored in a .csv file for now. A small example to illustrate my issue is as shown below:

source_node  target_node  ts  node_feats  edge_feats_1  edge_feats_2
            0            1   0         0.1           0.3           1.3
            0            3   0         0.3           1.3           2.3
            1            2   0         2.0           2.2           2.4
            2            3   0         0.4           1.4           5.4

As it can be noted, the dataset corresponds to a small graph with 4 nodes and 4 edges. Each node is endowed with a node feature and each edge carries two edges features.

The timestamps of the temporal edges are all set to 0 as this is the initial graph. The next set of data which is stored in another .csv contains edges with timestamps set to 1 as the edge features are set to evolve at every t += 1s.

• I understand the edge features are stored in the msg attribute, but then can we, also, simply assume that the node features simply there as well?

Here are my node features (volume) and edge features (surface_tension and L2_norm) (just a simple example based on the given small dataset above):

features = ['volume', 'surface_tension', 'L2_norm']
features_dict = {
    features_dict[0]: [1.0, 1.0, 1.0, 1.0],
    features_dict[1]: [0.7, 0.7, 0.7, 0.7],
    features_dict[2]: [0.8, 0.9, 1.7, 0.6]
}
feature_tensor_array = Tensor([feature_dict[i] for i in features])

# To check if the features are of the correct type and in correct form
# print(message_tensor_array[0].view(-1, 1))

I now proceed to define my temporal data:

import pandas as pd

df = pd.read_csv('geometry_utils/test_graph_data/raw/graph_data.csv')
remove_index = [''] * len(df)
df.index = remove_index
source = df['source_node']
destination = df['target_node']
timestamps = df['ts']

# Define the temporal data
data = TemporalData(
    src=LongTensor(list(source)),    # Source nodes involves in an event
    dst=LongTensor(list(destination)),   # Destination nodes involved in event
    t=LongTensor(list(timestamps)),  # Timestamps of events
    msg= feature_tensor_array[0].view(-1, 1),  # Storing edge features
)
data = data.to(device)

• As it can be noted I included only the first feature corresponding to the node feature and I am not how to proceed for the edge features? Should I consider implementing something like msg = [msg_0, msg_1, msg_2]? Should I then, also, alter something in the code. in tgn.py?

I have looked in temporal.py for further Information on msg but did not find much.

Thanks for any help in advance!

[rusty1s]

I understand the edge features are stored in the msg attribute, but then can we, also, simply assume that the node features simply there as well?

Yes, that's how I would do it as well. In this case, msg should be a [num_edges, 3] matrix, holding node_feats, edge_feats_1 and edge_feats_2.

[SalvishGoomanee]

ok got it, thanks!

[SalvishGoomanee]

In light of the above, my node features now includes the 3D position of the nodes and I was wondering if I should modify the tgn.py file to address the issues of rational and translational equivariance or if I can first input those in the messages at first and see what gives?

(f course this coordinate feature is a [num_nodes, 3] matrix)

[rusty1s]

I would pass the relative coordinates or the two absolute coordinates as messages TBH. This sounds like the easiest approach to integrate.

[SalvishGoomanee]

Ok thanks, so just adding them to node_feats essentially? (my node features file is a .csv with some scalar features for each node and consequently the x, y, z positions), i.e: node_id, x, y, z, feat_1, feat_2 etc...