/
node2vec_arxiv.py
60 lines (47 loc) · 2.25 KB
/
node2vec_arxiv.py
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
import argparse
import torch
from torch_geometric.nn import Node2Vec
from torch_geometric.utils import to_undirected
from ogb.nodeproppred import PygNodePropPredDataset
import os.path as osp
def save_embedding(model):
torch.save(model.embedding.weight.data.cpu(), 'ogbn/embedding_arxiv.pt')
def main():
parser = argparse.ArgumentParser(description='OGBN-Arxiv (Node2Vec)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--embedding_dim', type=int, default=128)
parser.add_argument('--walk_length', type=int, default=80)
parser.add_argument('--context_size', type=int, default=20)
parser.add_argument('--walks_per_node', type=int, default=10)
parser.add_argument('--batch_size', type=int, default=256)
parser.add_argument('--lr', type=float, default=0.01)
parser.add_argument('--epochs', type=int, default=5)
parser.add_argument('--log_steps', type=int, default=1)
args = parser.parse_args()
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)
# root = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'arxiv')
dataset = PygNodePropPredDataset(name='ogbn-arxiv',root = 'ogbn')
data = dataset[0]
data.edge_index = to_undirected(data.edge_index, data.num_nodes)
model = Node2Vec(data.edge_index, args.embedding_dim, args.walk_length,
args.context_size, args.walks_per_node,
sparse=True).to(device)
loader = model.loader(batch_size=args.batch_size, shuffle=True,
num_workers=4)
optimizer = torch.optim.SparseAdam(model.parameters(), lr=args.lr)
model.train()
for epoch in range(1, args.epochs + 1):
for i, (pos_rw, neg_rw) in enumerate(loader):
optimizer.zero_grad()
loss = model.loss(pos_rw.to(device), neg_rw.to(device))
loss.backward()
optimizer.step()
if (i + 1) % args.log_steps == 0:
print(f'Epoch: {epoch:02d}, Step: {i+1:03d}/{len(loader)}, '
f'Loss: {loss:.4f}')
if (i + 1) % 100 == 0: # Save model every 100 steps.
save_embedding(model)
save_embedding(model)
if __name__ == "__main__":
main()