JittorGeometric/examples/jodie_example.py at main · ChenJY-Count/JittorGeometric · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
import os.path as osp
import sys
root = osp.dirname(osp.dirname(osp.abspath(__file__)))
sys.path.append(root)
import jittor as jt
from jittor.nn import Linear
from jittor_geometric.datasets import JODIEDataset, TemporalDataLoader
from sklearn.metrics import average_precision_score, roc_auc_score
from tqdm import tqdm
import numpy as np
from jittor_geometric.datasets.tgb_seq import TGBSeqDataset
from jittor_geometric.data import TemporalData
from jittor_geometric.nn.models import JODIEEmbedding, compute_src_dst_node_time_shifts

# Use cuda
jt.flags.use_cuda = 1

# Load dataset from DGB or TGB-Seq
dataset_name = 'wikipedia' # wikipedia, mooc, reddit, lastfm
if dataset_name in [ 'wikipedia', 'reddit', 'mooc', 'lastfm']:
    # Load dataset from DGB
    path = osp.join(osp.dirname(osp.realpath(__file__)), 'data', 'JODIE')
    dataset = JODIEDataset(path, name=dataset_name)
    data = dataset[0]
    min_dst_idx, max_dst_idx = int(data.dst.min()), int(data.dst.max())

    # Split the dataset into train/val/test sets
    train_data, val_data, test_data = data.train_val_test_split(val_ratio=0.15, test_ratio=0.15)

    # Create TemporalDataLoader objects
    train_loader = TemporalDataLoader(train_data, batch_size=200, neg_sampling_ratio=1.0)
    val_loader = TemporalDataLoader(val_data, batch_size=200, neg_sampling_ratio=1.0)
    test_loader = TemporalDataLoader(test_data, batch_size=200, neg_sampling_ratio=1.0)
elif dataset_name in ['GoogleLocal', 'Yelp', 'Taobao', 'ML-20M' 'Flickr', 'YouTube', 'Patent', 'WikiLink']:
    # Load dataset from TGB-Seq
    path = osp.join(osp.dirname(osp.realpath(__file__)), 'data')
    dataset = TGBSeqDataset(root=path, name=dataset_name)
    train_idx=np.nonzero(dataset.train_mask)[0]
    val_idx=np.nonzero(dataset.val_mask)[0]
    test_idx=np.nonzero(dataset.test_mask)[0]
    edge_ids=np.arange(dataset.num_edges)+1
    if dataset.edge_feat is None:
        edge_feat=np.zeros((dataset.num_edges+1, 172))
    else:
        edge_feat=dataset.edge_feat
    if dataset.test_ns is not None:
        data = TemporalData(src=jt.array(dataset.src_node_ids.astype(np.int32)), dst=jt.array(dataset.dst_node_ids.astype(np.int32)), t=jt.array(dataset.time), msg=jt.array(edge_feat), train_mask=jt.array(train_idx.astype(np.int32)), val_mask=jt.array(val_idx.astype(np.int32)), test_mask=jt.array(test_idx.astype(np.int32)), test_ns=jt.array(dataset.test_ns.astype(np.int32)), edge_ids=jt.array(edge_ids.astype(np.int32)))
    else:
        data = TemporalData(src=jt.array(dataset.src_node_ids.astype(np.int32)), dst=jt.array(dataset.dst_node_ids.astype(np.int32)), t=jt.array(dataset.time), msg=jt.array(edge_feat), train_mask=jt.array(train_idx.astype(np.int32)), val_mask=jt.array(val_idx.astype(np.int32)), test_mask=jt.array(test_idx.astype(np.int32)), edge_ids=jt.array(edge_ids.astype(np.int32)))

    # Split the dataset into train/val/test sets
    train_data, val_data, test_data = data.train_val_test_split_w_mask()

    # Create TemporalDataLoader objects
    train_loader = TemporalDataLoader(train_data, batch_size=200, num_neg_sample=1)
    val_loader = TemporalDataLoader(val_data, batch_size=200, num_neg_sample=1)
    test_loader = TemporalDataLoader(test_data, batch_size=200, num_neg_sample=1)

# Compute time shift
src_node_mean_time_shift, src_node_std_time_shift, dst_node_mean_time_shift, dst_node_std_time_shift = compute_src_dst_node_time_shifts(
    src_node_ids=data.src.numpy(),
    dst_node_ids=data.dst.numpy(),
    node_interact_times=data.t.numpy()
)

# Define MLP-based predictor
class LinkPredictor(jt.nn.Module):
    def __init__(self, in_channels):
        super(LinkPredictor, self).__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 execute(self, z_src, z_dst):
        h = self.lin_src(z_src) + self.lin_dst(z_dst)
        h = jt.nn.relu(h)
        return self.lin_final(h)

# Define Memory module
embedding_dim = 10
num_users = int(data.src.max()) + 1
num_items = int(data.dst.max()) + 1
jodie_model = JODIEEmbedding(embedding_dim, num_users, num_items,
                                    src_node_mean_time_shift, src_node_std_time_shift,
                                    dst_node_mean_time_shift, dst_node_std_time_shift)

predictor = LinkPredictor(embedding_dim)
model = jt.nn.Sequential(jodie_model, predictor)
optimizer = jt.nn.Adam(list(model.parameters()),lr=0.0001)
criterion = jt.nn.BCEWithLogitsLoss()

def train():
    model.train()
    total_loss = 0
    for batch in tqdm(train_loader):
        user_idx = batch.src
        item_idx = batch.dst
        timestamp = batch.t
        neg_item_idx = batch.neg_dst

        # Get updated memory of all users and items involved in the computation
        pos_user_emb, pos_item_emb = model[0](user_idx, item_idx, timestamp)
        neg_user_emb, neg_item_emb = model[0](user_idx, neg_item_idx, timestamp)

        # Compute predictions and loss
        pos_pred = model[1](pos_user_emb, pos_item_emb)
        neg_pred = model[1](neg_user_emb, neg_item_emb)
        loss = criterion(pos_pred, jt.ones_like(pos_pred)) + criterion(neg_pred, jt.zeros_like(neg_pred))

        # Backpropagation and optimization
        optimizer.zero_grad()
        optimizer.step(loss)
        total_loss += float(loss) * batch.num_events

    return total_loss / train_data.num_events


def test(loader):
    model.eval()
    aps, aucs = [], []
    for batch in loader:
        user_idx = batch.src
        item_idx = batch.dst
        timestamp = batch.t
        neg_item_idx = jt.randint(0, num_items, (user_idx.shape[0],))

        # Get updated memory of all users and items involved in the computation
        pos_user_emb, pos_item_emb = model[0](user_idx, item_idx, timestamp)
        neg_user_emb, neg_item_emb = model[0](user_idx, neg_item_idx, timestamp)

        # Compute predictions
        pos_pred = model[1](pos_user_emb, pos_item_emb)
        neg_pred = model[1](neg_user_emb, neg_item_emb)
        y_pred = jt.concat([pos_pred.sigmoid(), neg_pred.sigmoid()], dim=0).numpy()
        y_true = jt.concat([jt.ones(pos_pred.shape[0]), jt.zeros(neg_pred.shape[0])], dim=0).numpy()

        # Compute metrics
        aps.append(average_precision_score(y_true, y_pred))
        aucs.append(roc_auc_score(y_true, y_pred))

    return np.mean(aps), np.mean(aucs)

best_ap = 0
patience = 5
save_model_path = osp.join(osp.dirname(osp.realpath(__file__)), 'data', 'saved_models')
for epoch in range(1, 11):
    loss = train()
    print(f'Epoch {epoch}, Loss: {loss:.4f}')
    val_ap, val_auc = test(val_loader)
    print(f'Val AP: {val_ap:.4f}, Val AUC: {val_auc:.4f}')
    if val_ap > best_ap and epoch >= 3:
        best_ap = val_ap
        # Save the model when achieving better performance on val set
        jt.save(model.state_dict(), f'{save_model_path}/{dataset_name}_model_jodie.pkl')
        print('Saved model is updated')
        patience = 5
    elif val_ap <= best_ap and epoch >= 3:
        patience -= 1
        # Early stop if patience decreases to zero
        if patience == 0:
            break

# Load the saved model for testing
model.load_state_dict(jt.load(f'{save_model_path}/{dataset_name}_model_jodie.pkl'))
test_ap, test_auc = test(test_loader)
print(f'Test AP: {test_ap:.4f}, Test AUC: {test_auc:.4f}')