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
|
import os.path as osp
import torch
from torch_geometric.datasets import BitcoinOTC
from torch_geometric.nn import SignedGCN
name = 'BitcoinOTC-1'
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', name)
dataset = BitcoinOTC(path, edge_window_size=1)
# Generate dataset.
pos_edge_indices, neg_edge_indices = [], []
for data in dataset:
pos_edge_indices.append(data.edge_index[:, data.edge_attr > 0])
neg_edge_indices.append(data.edge_index[:, data.edge_attr < 0])
if torch.cuda.is_available():
device = torch.device('cuda')
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = torch.device('mps')
else:
device = torch.device('cpu')
pos_edge_index = torch.cat(pos_edge_indices, dim=1).to(device)
neg_edge_index = torch.cat(neg_edge_indices, dim=1).to(device)
# Build and train model.
model = SignedGCN(64, 64, num_layers=2, lamb=5).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
train_pos_edge_index, test_pos_edge_index = model.split_edges(pos_edge_index)
train_neg_edge_index, test_neg_edge_index = model.split_edges(neg_edge_index)
x = model.create_spectral_features(train_pos_edge_index, train_neg_edge_index)
def train():
model.train()
optimizer.zero_grad()
z = model(x, train_pos_edge_index, train_neg_edge_index)
loss = model.loss(z, train_pos_edge_index, train_neg_edge_index)
loss.backward()
optimizer.step()
return loss.item()
def test():
model.eval()
with torch.no_grad():
z = model(x, train_pos_edge_index, train_neg_edge_index)
return model.test(z, test_pos_edge_index, test_neg_edge_index)
for epoch in range(101):
loss = train()
auc, f1 = test()
print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, AUC: {auc:.4f}, '
f'F1: {f1:.4f}')
|