图可以定义为 G = (V, E),其中 V 是一组节点,E 是一组边。 边是两个节点之间的连接,例如节点A和D有一条边。 另外,重要的是要注意图可以是有向图或无向图。 例如,下面的图是无向的,因为 A 与 D 连接,D 与 A 连接。还有一件事,图可以获取不同的节点属性以及边属性,但就我们的目的而言,今天并不重要。
from torch_geometric.nn import Node2Vecimport os.path as ospimport torchfrom torch_geometric.datasets import Planetoidfrom tqdm.notebook import tqdm dataset ='Cora' path = osp.join('.', 'data', dataset) dataset =Planetoid(path, dataset)# dowload or load the Cora dataset data = dataset[0] device = 'cuda' if torch.cuda.is_available() else 'cpu' # check if cuda is available to send the model and tensors to the GPU
model =Node2Vec(data.edge_index, embedding_dim=128, walk_length=20, context_size=10, walks_per_node=10, num_negative_samples=1, p=1, q=1, sparse=True).to(device) deftrain(): model.train() total_loss =0for pos_rw, neg_rw intqdm(loader): optimizer.zero_grad() loss = model.loss(pos_rw.to(device), neg_rw.to(device)) loss.backward() optimizer.step() total_loss += loss.item()return total_loss /len(loader) for epoch inrange(1, 100): loss =train()print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}') all_vectors =""for tensor inmodel(torch.arange(data.num_nodes, device=device)): s ="\t".join([str(value) for value in tensor.detach().cpu().numpy()]) all_vectors += s +"\n" withopen("vectors.txt", "w")as f: f.write(all_vectors) withopen("labels.txt", "w")as f: f.write("\n".join([str(label) for label in data.y.numpy()]))