JODIE wikipedia

[xavierallem]

I have read the paper for JODIE, I have also gone through the pytorch geometric implementation. I understood the dataset. However if I want to predict the state change ie. Banned or Not banned. How can I? is there a Implementation already done?

[rusty1s]

Sorry, what do you mean by state change? Are you referring to node classification?

[xavierallem]

Yes Node classification, In the paper they have described it as state change.
is the node classification example code available? if not how can In realise it?

[EdisonLeeeee]

Sorry, there is no node classification example yet; will implement it soon.

[xavierallem]

@EdisonLeeeee I can try to implement it, Can you give me some pointers?

[EdisonLeeeee]

Given the input (src, dst, t, msg), you can implement it as an edge classification task, classifying the edge (src, dst, msg) based on the class assigned to src.

[xavierallem]

@EdisonLeeeee So as an example of edge classifier. I have Implemented it like using the y label given in the csv.

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

    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)

num_classes=len(torch.unique(data.y))
node_cls = NodeClassifier(in_channels=embedding_dim,num_classes=num_classes).to(device)
optimizer = torch.optim.Adam(
    set(memory.parameters()) | set(gnn.parameters())
    | set(node_cls.parameters()), lr=0.0001)
criterion = torch.nn.CrossEntropyLoss()

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

Train Code

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

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

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

        n_id, edge_index, e_id = neighbor_loader(batch.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 = node_cls(z[assoc[batch.src]], z[assoc[batch.dst]])
        # Debugging shapes
      


        loss = criterion(pos_out, batch.y) #here

        # Update memory and neighbor loader with ground-truth state.
        memory.update_state(batch.src, batch.dst, batch.t, batch.msg)
        neighbor_loader.insert(batch.src, batch.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()
    node_cls.eval()

    torch.manual_seed(12345)  # Ensure deterministic sampling across epochs.
    correct,total=0,0

    ys, preds = [], []
    for batch in loader:
        batch = batch.to(device)

        n_id, edge_index, e_id = neighbor_loader(batch.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))
        out = node_cls(z[assoc[batch.src]],z[assoc[batch.dst]])

        pred=out.argmax(dim=-1) #here 
        correct+=(pred==batch.y).sum().item()
        total+=batch.y.size(0)

        memory.update_state(batch.src, batch.dst, batch.t, batch.msg)
        neighbor_loader.insert(batch.src, batch.dst)

 


    return correct/total
for epoch in range(1, 51):
    loss = train()
    print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}')
    val_auc = test(val_loader)
    test_auc = test(test_loader)
    print(f'Val Accuracy: {val_auc:.4f}')
    print(f'Test Accuracy: {test_auc:.4f}')

Some changes I did, Is this the correct way for edge classifier?

[xavierallem]

@rusty1s Hi, Any update regarding this?

[rusty1s]

Hey, sorry. I am not sure I understand why the code above corresponds a node-level task since you use both src and dst features (i.e., a pair of nodes/an edge) to do the prediction.

[xavierallem]

Hey, sorry. I am not sure I understand why the code above corresponds a node-level task since you use both src and dst features (i.e., a pair of nodes/an edge) to do the prediction.

IHey @rusty1s , I was trying to do Edge Classification as suggested by a peer instead of node
However I have updated the code. I have assigned classes randomly according to the message feature, I am getting way good accuracy which i feel is wrong and stable. I am relatively new to temporal pytorch implementation this way, sorry for the little information I have. I have tried to use the node classfication from example and tried to merge it with this

from collections import Counter
def assign_class(value):
   if value < -1.0:
       return 0
   elif -1.0 <= value < 0.0:
       return 1
   elif 0.0 <= value < 1.0:
       return 2
   elif 1.0 <= value < 2.0:
       return 3
   elif 2.0 <= value < 3.0:
       return 4
   elif 3.0 <= value < 4.0:
       return 5
   else:
       return 6



data.y = data.msg.sum(dim=1)
data.y = torch.tensor([assign_class(val.item()) for val in data.y], device='cuda:0')

print(data.y)
class_counts = Counter(data.y.cpu().numpy())

# Print the class counts
for cls, count in class_counts.items():
   print(f'Class {cls}: {count} members')

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 NodeClassifier(torch.nn.Module):
   def __init__(self, in_channels):
       super().__init__()
       self.lin_dst = Linear(in_channels, in_channels)
       self.lin_final = Linear(in_channels, num_classes)

       



   def forward(self, x):
       x = self.lin_dst(x)
       x = x.relu()
       x = F.dropout(x, p=0.5, training=self.training)
       x = self.lin_final(x)
       return x

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

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

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

       n_id, edge_index, e_id = neighbor_loader(batch.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 = node_cls(z[assoc[batch.src]])
      
      # print(prob)
       #print(y)
     
       loss = criterion(pos_out, batch.y)


       memory.update_state(batch.src, batch.dst, batch.t, batch.msg)
       neighbor_loader.insert(batch.src, batch.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()
   node_cls.eval()

   torch.manual_seed(12345)  # Ensure deterministic sampling across epochs.
   correct,total=0,0

   ys, preds = [], []
   for batch in loader:
       batch = batch.to(device)

       n_id, edge_index, e_id = neighbor_loader(batch.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))


       out=node_cls(z[assoc[batch.src]])
       

       pred=out.argmax(dim=-1)
       correct+=((pred==batch.y).sum().item())
       total+=batch.y.size(0)

       memory.update_state(batch.src, batch.dst, batch.t, batch.msg)
       neighbor_loader.insert(batch.src, batch.dst)



   return correct/total
for epoch in range(1, 51):
   loss = train()
   print(f'Epoch: {epoch:02d}, Loss: {loss:.8f}')
   val_auc = test(val_loader)
   test_auc = test(test_loader)
   print(f'Val Accuracy: {val_auc:.8f}')
   print(f'Test Accuracy: {test_auc:.8f}')

[rusty1s]

Your code looks alright to me. I am curious why you do

node_cls(z[assoc[batch.src]])

instead of

link_cls(z[assoc[batch.src]], z[assoc[batch.dst]])

though.