├── .gitignore ├── README.md ├── beginner ├── __init__.py ├── cifar10_tutorial.py ├── neural_networks_tutorial.py └── two_layer_net.py ├── dlwizard ├── __init__.py ├── cnn.py ├── common.py ├── fnn.py ├── linear_regression.py ├── logistic_regression.py ├── lstm.py └── rnn.py ├── gnn ├── __init__.py ├── cs │ ├── __init__.py │ ├── model.py │ ├── readme.md │ └── train.py ├── data │ ├── __init__.py │ ├── acm.py │ ├── aminer.py │ ├── dblp.py │ ├── dgl.py │ ├── heco.py │ └── imdb.py ├── dgl │ ├── __init__.py │ ├── dgl_first_demo.py │ ├── edge_clf.py │ ├── edge_clf_hetero.py │ ├── edge_clf_hetero_mb.py │ ├── edge_clf_mb.py │ ├── edge_type_hetero.py │ ├── graph_clf.py │ ├── graph_clf_hetero.py │ ├── link_pred.py │ ├── link_pred_hetero.py │ ├── link_pred_hetero_mb.py │ ├── link_pred_mb.py │ ├── model.py │ ├── node_clf.py │ ├── node_clf_hetero.py │ ├── node_clf_hetero_mb.py │ └── node_clf_mb.py ├── gat │ ├── __init__.py │ ├── model.py │ ├── readme.txt │ ├── train_inductive.py │ └── train_transductive.py ├── gcn │ ├── __init__.py │ ├── model.py │ ├── readme.txt │ └── train.py ├── han │ ├── __init__.py │ ├── model.py │ ├── readme.txt │ └── train.py ├── heco │ ├── __init__.py │ ├── model.py │ ├── readme.md │ └── train.py ├── hetgnn │ ├── __init__.py │ ├── eval.py │ ├── model.py │ ├── preprocess.py │ ├── random_walk.py │ ├── readme.md │ ├── train.py │ └── utils.py ├── hgconv │ ├── __init__.py │ ├── model.py │ ├── readme.txt │ └── train.py ├── hgt │ ├── __init__.py │ ├── model.py │ ├── readme.txt │ └── train.py ├── lp │ ├── __init__.py │ ├── model.py │ ├── readme.txt │ └── train.py ├── magnn │ ├── __init__.py │ ├── encoder.py │ ├── model.py │ ├── readme.txt │ ├── train_dblp.py │ └── train_imdb.py ├── metapath2vec │ ├── __init__.py │ ├── random_walk.py │ ├── readme.txt │ ├── skipgram.py │ ├── train.py │ └── train_word2vec.py ├── rgcn │ ├── __init__.py │ ├── model.py │ ├── model_hetero.py │ ├── readme.txt │ ├── train_entity_clf.py │ └── train_link_pred.py ├── rhgnn │ ├── __init__.py │ ├── model.py │ ├── readme.txt │ └── train.py ├── sign │ ├── __init__.py │ ├── model.py │ ├── readme.txt │ └── train.py ├── supergat │ ├── __init__.py │ ├── attention.py │ ├── model.py │ ├── readme.txt │ └── train.py └── utils │ ├── __init__.py │ ├── data.py │ ├── metapath.py │ ├── metrics.py │ ├── neg_sampler.py │ └── random_walk.py ├── kgrec ├── __init__.py └── kgcn │ ├── __init__.py │ ├── data.py │ ├── dataloader.py │ ├── model.py │ ├── readme.txt │ └── train.py ├── nlp ├── __init__.py └── tfms │ ├── __init__.py │ ├── causal_lm_model.py │ ├── eqa_model.py │ ├── eqa_pipeline.py │ ├── masked_lm_model.py │ ├── masked_lm_pipeline.py │ ├── seq_clf_model.py │ ├── seq_clf_pipeline.py │ ├── text_gen_model.py │ └── text_gen_pipeline.py └── requirements.txt /.gitignore: -------------------------------------------------------------------------------- 1 | # PyCharm 2 | /.idea/ 3 | 4 | # Python 5 | __pycache__/ 6 | 7 | # Dataset 8 | /data/ 9 | -------------------------------------------------------------------------------- /beginner/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/beginner/__init__.py -------------------------------------------------------------------------------- /beginner/cifar10_tutorial.py: -------------------------------------------------------------------------------- 1 | """参考：""" 2 | import matplotlib.pyplot as plt 3 | import numpy as np 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | import torch.optim as optim 8 | import torchvision 9 | import torchvision.transforms as transforms 10 | from torch.utils.data import DataLoader 11 | 12 | classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') 13 | PATH = 'cifar_net.pth' 14 | 15 | 16 | def imshow(img): 17 | img = img / 2 + 0.5 # unnormalize 18 | npimg = img.numpy() 19 | plt.imshow(np.transpose(npimg, (1, 2, 0))) 20 | plt.show() 21 | 22 | 23 | def show_random_image(trainloader): 24 | # get some random training images 25 | dataiter = iter(trainloader) 26 | images, labels = next(dataiter) 27 | 28 | # show images 29 | imshow(torchvision.utils.make_grid(images)) 30 | # print labels 31 | print(' '.join('{:>5}'.format(classes[labels[j]]) for j in range(4))) 32 | 33 | 34 | class Net(nn.Module): 35 | 36 | def __init__(self): 37 | super().__init__() 38 | self.conv1 = nn.Conv2d(3, 6, 5) 39 | self.pool = nn.MaxPool2d(2, 2) 40 | self.conv2 = nn.Conv2d(6, 16, 5) 41 | self.fc1 = nn.Linear(16 * 5 * 5, 120) 42 | self.fc2 = nn.Linear(120, 84) 43 | self.fc3 = nn.Linear(84, 10) 44 | 45 | def forward(self, x): 46 | x = self.pool(F.relu(self.conv1(x))) 47 | x = self.pool(F.relu(self.conv2(x))) 48 | x = x.view(-1, 16 * 5 * 5) 49 | x = F.relu(self.fc1(x)) 50 | x = F.relu(self.fc2(x)) 51 | x = self.fc3(x) 52 | return x 53 | 54 | 55 | def train(trainloader): 56 | device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 57 | print(f'device = {device}') 58 | net = Net() 59 | net.to(device) 60 | 61 | criterion = nn.CrossEntropyLoss() 62 | optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) 63 | for epoch in range(2): # loop over the dataset multiple times 64 | running_loss = 0.0 65 | for i, data in enumerate(trainloader): 66 | # get the inputs; data is a list of [inputs, labels] 67 | inputs, labels = data[0].to(device), data[1].to(device) 68 | 69 | # zero the parameter gradients 70 | optimizer.zero_grad() 71 | 72 | # forward + backward + optimize 73 | outputs = net(inputs) 74 | loss = criterion(outputs, labels) 75 | loss.backward() 76 | optimizer.step() 77 | 78 | # print statistics 79 | running_loss += loss.item() 80 | if i % 2000 == 1999: # print every 2000 mini-batches 81 | print('[{:d}, {:5d}] loss: {:.3f}'.format(epoch + 1, i + 1, running_loss / 2000)) 82 | running_loss = 0.0 83 | print('Finished Training') 84 | torch.save(net.state_dict(), PATH) 85 | 86 | 87 | def test(testloader): 88 | net = Net() 89 | net.load_state_dict(torch.load(PATH)) 90 | 91 | correct = 0 92 | total = 0 93 | with torch.no_grad(): 94 | for data in testloader: 95 | images, labels = data 96 | outputs = net(images) 97 | _, predicted = torch.max(outputs.data, 1) 98 | total += labels.size(0) 99 | correct += (predicted == labels).sum().item() 100 | print('Accuracy of the network on the 10000 test images: {:.2%}'.format(correct / total)) 101 | 102 | 103 | def class_accuracy(testloader): 104 | net = Net() 105 | net.load_state_dict(torch.load(PATH)) 106 | 107 | class_correct = [0.] * 10 108 | class_total = [0.] * 10 109 | with torch.no_grad(): 110 | for data in testloader: 111 | images, labels = data 112 | outputs = net(images) 113 | _, predicted = torch.max(outputs, 1) 114 | c = (predicted == labels).squeeze() 115 | for i in range(4): 116 | label = labels[i] 117 | class_correct[label] += c[i].item() 118 | class_total[label] += 1 119 | 120 | for i in range(10): 121 | print('Accuracy of {:>5} : {:.2%}'.format(classes[i], class_correct[i] / class_total[i])) 122 | 123 | 124 | def main(): 125 | transform = transforms.Compose([ 126 | transforms.ToTensor(), 127 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) 128 | ]) 129 | 130 | trainset = torchvision.datasets.CIFAR10('D:\\torchdata', True, transform, download=True) 131 | trainloader = DataLoader(trainset, batch_size=4, shuffle=True, num_workers=0) 132 | testset = torchvision.datasets.CIFAR10('D:\\torchdata', False, transform, download=True) 133 | testloader = DataLoader(testset, batch_size=4, shuffle=True, num_workers=0) 134 | # show_random_image(trainloader) 135 | 136 | train(trainloader) 137 | test(testloader) 138 | class_accuracy(testloader) 139 | 140 | 141 | if __name__ == '__main__': 142 | main() 143 | -------------------------------------------------------------------------------- /beginner/neural_networks_tutorial.py: -------------------------------------------------------------------------------- 1 | """参考：""" 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class Net(nn.Module): 7 | 8 | def __init__(self): 9 | super().__init__() 10 | # 1 input image channel, 6 output channels, 3x3 square convolution 11 | # kernel 12 | self.conv1 = nn.Conv2d(1, 6, 3) 13 | self.conv2 = nn.Conv2d(6, 16, 3) 14 | # an affine operation: y = Wx + b 15 | self.fc1 = nn.Linear(16 * 6 * 6, 120) # 6*6 from image dimension 16 | self.fc2 = nn.Linear(120, 84) 17 | self.fc3 = nn.Linear(84, 10) 18 | 19 | def forward(self, x): 20 | # Max pooling over a (2, 2) window 21 | x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2)) 22 | # If the size is a square you can only specify a single number 23 | x = F.max_pool2d(F.relu(self.conv2(x)), 2) 24 | x = x.view(-1, self.num_flat_features(x)) 25 | x = F.relu(self.fc1(x)) 26 | x = F.relu(self.fc2(x)) 27 | x = self.fc3(x) 28 | return x 29 | 30 | def num_flat_features(self, x): 31 | size = x.size()[1:] # all dimensions except the batch dimension 32 | num_features = 1 33 | for s in size: 34 | num_features *= s 35 | return num_features 36 | 37 | 38 | if __name__ == '__main__': 39 | net = Net() 40 | print(net) 41 | -------------------------------------------------------------------------------- /beginner/two_layer_net.py: -------------------------------------------------------------------------------- 1 | """参考：""" 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import torch.optim as optim 6 | 7 | 8 | class TwoLayerNet(nn.Module): 9 | """包含两个全连接层的前馈神经网络""" 10 | 11 | def __init__(self, D_in, H, D_out): 12 | """ 13 | In the constructor we instantiate two nn.Linear modules and assign them as member variables. 14 | """ 15 | super().__init__() 16 | self.linear1 = nn.Linear(D_in, H) 17 | self.linear2 = nn.Linear(H, D_out) 18 | 19 | def forward(self, x): 20 | """ 21 | In the forward function we accept a Tensor of input data and we must return 22 | a Tensor of output data. We can use Modules defined in the constructor as 23 | well as arbitrary operators on Tensors. 24 | """ 25 | return self.linear2(F.relu(self.linear1(x))) 26 | 27 | 28 | # N is batch size; D_in is input dimension; 29 | # H is hidden dimension; D_out is output dimension. 30 | N, D_in, H, D_out = 64, 1000, 100, 10 31 | 32 | # Create random Tensors to hold inputs and outputs 33 | x = torch.randn(N, D_in) 34 | y = torch.randn(N, D_out) 35 | 36 | # Construct our model by instantiating the class defined above 37 | model = TwoLayerNet(D_in, H, D_out) 38 | 39 | # Construct our loss function and an Optimizer. The call to model.parameters() 40 | # in the SGD constructor will contain the learnable parameters of the two 41 | # nn.Linear modules which are members of the model. 42 | criterion = nn.MSELoss(reduction='sum') 43 | optimizer = optim.SGD(model.parameters(), lr=1e-4) 44 | for t in range(500): 45 | # Forward pass: compute predicted y by passing x to the model. 46 | y_pred = model(x) 47 | 48 | # Compute and print loss. 49 | loss = criterion(y_pred, y) 50 | if t % 100 == 99: 51 | print(t, loss.item()) 52 | 53 | # Zero gradients, perform a backward pass, and update the weights. 54 | optimizer.zero_grad() 55 | loss.backward() 56 | optimizer.step() 57 | -------------------------------------------------------------------------------- /dlwizard/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/dlwizard/__init__.py -------------------------------------------------------------------------------- /dlwizard/cnn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | import torch.optim as optim 4 | 5 | from dlwizard.common import train_mnist 6 | 7 | 8 | # https://www.deeplearningwizard.com/deep_learning/practical_pytorch/pytorch_convolutional_neuralnetwork/ 9 | class CNN(nn.Module): 10 | 11 | def __init__(self): 12 | super().__init__() 13 | self.conv1 = nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2) 14 | self.pool1 = nn.MaxPool2d(kernel_size=2) 15 | self.conv2 = nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2) 16 | self.pool2 = nn.MaxPool2d(kernel_size=2) 17 | self.fc1 = nn.Linear(32 * 7 * 7, 10) 18 | 19 | def forward(self, x): 20 | """ 21 | :param x: tensor(*, 28, 28) 输入图像 22 | :return: tensor(*, 10) 23 | """ 24 | out = F.relu(self.conv1(x)) # (*, 16, 28, 28) 25 | out = self.pool1(out) # (*, 16, 14, 14) 26 | out = F.relu(self.conv2(out)) # (*, 32, 14, 14) 27 | out = self.pool2(out) # (*, 32, 7, 7) 28 | out = out.view(out.shape[0], -1) # (*, 32 * 7 * 7) 29 | out = self.fc1(out) # (*, 10) 30 | return out 31 | 32 | 33 | def main(): 34 | model = CNN() 35 | optimizer = optim.SGD(model.parameters(), lr=0.01) 36 | for p in model.parameters(): 37 | print(p.size()) 38 | train_mnist(model, optimizer, (-1, 1, 28, 28)) 39 | 40 | 41 | if __name__ == '__main__': 42 | main() 43 | -------------------------------------------------------------------------------- /dlwizard/common.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torchvision.transforms as transforms 4 | from torch.utils.data.dataloader import DataLoader 5 | from torchvision.datasets import MNIST 6 | 7 | 8 | def train_mnist(model, optimizer, view_shape=(-1, 28, 28)): 9 | train_dataset = MNIST('D:\\torchdata', True, transforms.ToTensor(), download=True) 10 | test_dataset = MNIST('D:\\torchdata', False, transforms.ToTensor()) 11 | batch_size = 100 12 | n_epochs = 5 13 | train_loader = DataLoader(train_dataset, batch_size, shuffle=True) 14 | test_loader = DataLoader(test_dataset, batch_size, shuffle=False) 15 | 16 | criterion = nn.CrossEntropyLoss() 17 | iteration = 0 18 | for epoch in range(n_epochs): 19 | for images, labels in train_loader: 20 | images = images.view(view_shape).requires_grad_() 21 | optimizer.zero_grad() 22 | outputs = model(images) 23 | loss = criterion(outputs, labels) 24 | loss.backward() 25 | optimizer.step() 26 | 27 | iteration += 1 28 | if iteration % 500 == 0: 29 | # Calculate Accuracy 30 | correct = 0 31 | for test_images, test_labels in test_loader: 32 | test_images = test_images.view(view_shape).requires_grad_() 33 | outputs = model(test_images) 34 | predicted = torch.argmax(outputs.data, dim=1) 35 | correct += (predicted == test_labels).sum().item() 36 | accuracy = correct / len(test_dataset) 37 | print('Iteration: {}. Loss: {}. Accuracy: {:.2%}'.format( 38 | iteration, loss.item(), accuracy 39 | )) 40 | -------------------------------------------------------------------------------- /dlwizard/fnn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.optim as optim 3 | 4 | from dlwizard.common import train_mnist 5 | 6 | 7 | # https://www.deeplearningwizard.com/deep_learning/practical_pytorch/pytorch_feedforward_neuralnetwork/ 8 | class FeedforwardNeuralNetwork(nn.Module): 9 | 10 | def __init__(self, input_dim, hidden_dim, output_dim, activation): 11 | super().__init__() 12 | self.fc1 = nn.Linear(input_dim, hidden_dim) 13 | self.activation = activation 14 | self.fc2 = nn.Linear(hidden_dim, output_dim) 15 | 16 | def forward(self, x): 17 | out = self.fc1(x) 18 | out = self.activation(out) 19 | out = self.fc2(out) 20 | return out 21 | 22 | 23 | def main(): 24 | model = FeedforwardNeuralNetwork(28 * 28, 100, 10, nn.ReLU()) 25 | optimizer = optim.SGD(model.parameters(), lr=0.1) 26 | for p in model.parameters(): 27 | print(p.size()) 28 | train_mnist(model, optimizer, (-1, 28 * 28)) 29 | 30 | 31 | if __name__ == '__main__': 32 | main() 33 | -------------------------------------------------------------------------------- /dlwizard/linear_regression.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import numpy as np 3 | import torch 4 | import torch.nn as nn 5 | import torch.optim as optim 6 | 7 | 8 | # https://www.deeplearningwizard.com/deep_learning/practical_pytorch/pytorch_linear_regression/ 9 | class LinearRegressionModel(nn.Module): 10 | 11 | def __init__(self, input_dim, output_dim): 12 | super().__init__() 13 | self.linear = nn.Linear(input_dim, output_dim) 14 | 15 | def forward(self, x): 16 | return self.linear(x) 17 | 18 | 19 | def main(): 20 | X_train = np.arange(11, dtype=np.float32).reshape((-1, 1)) 21 | y_train = 2 * X_train + 1 22 | model = LinearRegressionModel(1, 1) 23 | criterion = nn.MSELoss() 24 | optimizer = optim.SGD(model.parameters(), lr=0.01) 25 | epochs = 100 26 | for epoch in range(epochs): 27 | # Convert numpy array to torch tensor 28 | inputs = torch.from_numpy(X_train).requires_grad_() 29 | labels = torch.from_numpy(y_train) 30 | 31 | # Clear gradients w.r.t. parameters 32 | optimizer.zero_grad() 33 | 34 | # Forward to get output 35 | outputs = model(inputs) 36 | 37 | # Calculate loss 38 | loss = criterion(outputs, labels) 39 | 40 | # Getting gradients w.r.t. parameters 41 | loss.backward() 42 | 43 | # Updating parameters 44 | optimizer.step() 45 | 46 | print('epoch {}, loss {}'.format(epoch, loss.item())) 47 | 48 | # Purely inference 49 | predicted = model(torch.from_numpy(X_train).requires_grad_()).data.numpy() 50 | print(predicted) 51 | 52 | # Clear figure 53 | plt.clf() 54 | 55 | # Plot true data 56 | plt.plot(X_train, y_train, 'go', label='True data', alpha=0.5) 57 | 58 | # Plot predictions 59 | plt.plot(X_train, predicted, '--', label='Predictions', alpha=0.5) 60 | 61 | # Legend and plot 62 | plt.legend(loc='best') 63 | plt.show() 64 | 65 | 66 | if __name__ == '__main__': 67 | main() 68 | -------------------------------------------------------------------------------- /dlwizard/logistic_regression.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.optim as optim 3 | 4 | from dlwizard.common import train_mnist 5 | 6 | 7 | # https://www.deeplearningwizard.com/deep_learning/practical_pytorch/pytorch_logistic_regression/ 8 | class LogisticRegressionModel(nn.Module): 9 | 10 | def __init__(self, input_dim, output_dim): 11 | super().__init__() 12 | self.linear = nn.Linear(input_dim, output_dim) 13 | 14 | def forward(self, x): 15 | return self.linear(x) 16 | 17 | 18 | def main(): 19 | model = LogisticRegressionModel(28 * 28, 10) 20 | optimizer = optim.SGD(model.parameters(), lr=0.001) 21 | for p in model.parameters(): 22 | print(p.size()) 23 | train_mnist(model, optimizer, (-1, 28 * 28)) 24 | 25 | 26 | if __name__ == '__main__': 27 | main() 28 | -------------------------------------------------------------------------------- /dlwizard/lstm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.optim as optim 4 | 5 | from dlwizard.common import train_mnist 6 | 7 | 8 | # https://www.deeplearningwizard.com/deep_learning/practical_pytorch/pytorch_lstm_neuralnetwork/ 9 | class LSTM(nn.Module): 10 | 11 | def __init__(self, in_dim, hidden_dim, num_layers, out_dim): 12 | super().__init__() 13 | self.hidden_dim = hidden_dim 14 | self.num_layers = num_layers 15 | self.lstm = nn.LSTM(in_dim, hidden_dim, num_layers, batch_first=True) 16 | self.fc = nn.Linear(hidden_dim, out_dim) 17 | 18 | def forward(self, x): 19 | """ 20 | :param x: tensor(batch, seq_len, d_in) 21 | :return: tensor(batch, d_out) 22 | """ 23 | # 由于设置了batch_first=True，因此x是(batch, seq_len, d_in)而不是(seq_len, batch, d_in) 24 | # Initialize hidden state with zeros (num_layers, batch, d_hid) 25 | h0 = torch.zeros(self.num_layers, x.shape[0], self.hidden_dim).requires_grad_() 26 | 27 | # Initialize cell state (num_layers, batch, d_hid) 28 | c0 = torch.zeros(self.num_layers, x.shape[0], self.hidden_dim).requires_grad_() 29 | 30 | # 28 time steps 31 | # We need to detach as we are doing truncated backpropagation through time (BPTT) 32 | # Otherwise we'll backprop all the way to the start even after going through another batch 33 | # out: (batch, seq_len, d_hid), hn, cn: (num_layers, batch, d_hid) 34 | out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach())) 35 | 36 | # Index hidden state of last time step 37 | # out[:, -1, :] -> (batch, d_hid) --> just want last time step hidden states! 38 | out = self.fc(out[:, -1, :]) # (batch, d_out) 39 | return out 40 | 41 | 42 | def main(): 43 | in_dim = 28 44 | seq_len = 28 45 | model = LSTM(in_dim, 100, 1, 10) 46 | optimizer = optim.SGD(model.parameters(), lr=0.1) 47 | for p in model.parameters(): 48 | print(p.size()) 49 | train_mnist(model, optimizer, (-1, seq_len, in_dim)) 50 | 51 | 52 | if __name__ == '__main__': 53 | main() 54 | -------------------------------------------------------------------------------- /dlwizard/rnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.optim as optim 4 | 5 | from dlwizard.common import train_mnist 6 | 7 | 8 | # https://www.deeplearningwizard.com/deep_learning/practical_pytorch/pytorch_recurrent_neuralnetwork/ 9 | class RNN(nn.Module): 10 | 11 | def __init__(self, in_dim, hidden_dim, num_layers, out_dim): 12 | super().__init__() 13 | self.hidden_dim = hidden_dim 14 | self.num_layers = num_layers 15 | self.rnn = nn.RNN(in_dim, hidden_dim, num_layers, nonlinearity='tanh', batch_first=True) 16 | self.fc = nn.Linear(hidden_dim, out_dim) 17 | 18 | def forward(self, x): 19 | """ 20 | :param x: tensor(batch, seq_len, d_in) 21 | :return: tensor(batch, d_out) 22 | """ 23 | # 由于设置了batch_first=True，因此x是(batch, seq_len, d_in)而不是(seq_len, batch, d_in) 24 | # Initialize hidden state with zeros (num_layers, batch, d_hid) 25 | h0 = torch.zeros(self.num_layers, x.shape[0], self.hidden_dim).requires_grad_() 26 | 27 | # We need to detach the hidden state to prevent exploding/vanishing gradients 28 | # This is part of truncated backpropagation through time (BPTT) 29 | # out: (batch, seq_len, d_hid), hn: (num_layers, batch, d_hid) 30 | out, hn = self.rnn(x, h0.detach()) 31 | 32 | # Index hidden state of last time step 33 | # out[:, -1, :] -> (batch, d_hid) --> just want last time step hidden states! 34 | out = self.fc(out[:, -1, :]) # (batch, d_out) 35 | return out 36 | 37 | 38 | def main(): 39 | in_dim = 28 40 | seq_len = 28 41 | model = RNN(in_dim, 100, 2, 10) 42 | optimizer = optim.SGD(model.parameters(), lr=0.1) 43 | for p in model.parameters(): 44 | print(p.size()) 45 | train_mnist(model, optimizer, (-1, seq_len, in_dim)) 46 | 47 | 48 | if __name__ == '__main__': 49 | main() 50 | -------------------------------------------------------------------------------- /gnn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/__init__.py -------------------------------------------------------------------------------- /gnn/cs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/cs/__init__.py -------------------------------------------------------------------------------- /gnn/cs/readme.md: -------------------------------------------------------------------------------- 1 | # Cora 2 | ## Linear + C&S 3 | `python -m gnn.cs.train --dataset=cora --base-model=Linear --num-hidden=64 --epochs=10` 4 | ``` 5 | Base model Linear 6 | Test Acc 0.4710 7 | C&S 8 | Test Acc 0.8010 9 | ``` 10 | 11 | ## MLP + C&S 12 | `python -m gnn.cs.train --dataset=cora --base-model=MLP --num-hidden=64 --epochs=10` 13 | ``` 14 | Base model MLP 15 | Test Acc 0.3190 16 | C&S 17 | Test Acc 0.7960 18 | ``` 19 | 20 | # Citeseer 21 | ## Linear + C&S 22 | `python -m gnn.cs.train --dataset=citeseer --base-model=Linear --num-hidden=64 --epochs=10` 23 | ``` 24 | Base model Linear 25 | Test Acc 0.4710 26 | C&S 27 | Test Acc 0.6650 28 | ``` 29 | 30 | ## MLP + C&S 31 | `python -m gnn.cs.train --dataset=citeseer --base-model=MLP --num-hidden=64 --epochs=10` 32 | ``` 33 | Base model MLP 34 | Test Acc 0.2310 35 | C&S 36 | Test Acc 0.6110 37 | ``` 38 | 39 | # Pubmed 40 | ## Linear + C&S 41 | `python -m gnn.cs.train --dataset=pubmed --base-model=Linear --num-hidden=64 --epochs=10` 42 | ``` 43 | Base model Linear 44 | Test Acc 0.6870 45 | C&S 46 | Test Acc 0.7780 47 | ``` 48 | 49 | ## MLP + C&S 50 | `python -m gnn.cs.train --dataset=pubmed --base-model=MLP --num-hidden=64 --epochs=10` 51 | ``` 52 | Base model MLP 53 | Test Acc 0.4350 54 | C&S 55 | Test Acc 0.7380 56 | ``` 57 | 58 | # ogbn-products 59 | ## Linear + C&S 60 | `python -m gnn.cs.train --dataset=ogbn-products --ogb-root=/home/zzy/ogb --base-model=Linear --correct-alpha=0.6 --smooth-alpha=0.9` 61 | ``` 62 | Base model Linear 63 | Test Acc 0.4777 64 | C&S 65 | Test Acc 0.7220 66 | ``` 67 | 68 | # ogbn-arxiv 69 | ## Linear + C&S 70 | `python -m gnn.cs.train --dataset=ogbn-arxiv --ogb-root=/home/zzy/ogb --base-model=Linear --correct-alpha=0.8 --correct-norm=right --smooth-alpha=0.6` 71 | ``` 72 | Base model Linear 73 | Test Acc 0.5245 74 | C&S 75 | Test Acc 0.6776 76 | ``` 77 | 78 | ## MLP + C&S 79 | `python -m gnn.cs.train --dataset=ogbn-arxiv --ogb-root=/home/zzy/ogb --base-model=MLP --correct-alpha=0.979 --correct-norm=left --smooth-alpha=0.756 --smooth-norm=right` 80 | ``` 81 | Base model MLP 82 | Test Acc 0.5616 83 | C&S 84 | Test Acc 0.6717 85 | ``` 86 | -------------------------------------------------------------------------------- /gnn/data/__init__.py: -------------------------------------------------------------------------------- 1 | from .dgl import UserItemDataset, RandomGraphDataset 2 | from .acm import ACMDataset, ACM3025Dataset 3 | from .dblp import DBLPFourAreaDataset, DBLP4057Dataset 4 | from .imdb import IMDbDataset, IMDb5kDataset 5 | from .aminer import AMinerCSDataset, AMiner2Dataset 6 | from .heco import ACMHeCoDataset, DBLPHeCoDataset, FreebaseHeCoDataset, AMinerHeCoDataset 7 | -------------------------------------------------------------------------------- /gnn/data/dgl.py: -------------------------------------------------------------------------------- 1 | import dgl 2 | import numpy as np 3 | import torch 4 | from dgl.data import DGLDataset 5 | 6 | 7 | class UserItemDataset(DGLDataset): 8 | 9 | def __init__(self, n_users=1000, n_items=500, n_follows=3000, n_clicks=5000, n_dislikes=500, 10 | n_features=10, n_user_classes=5, n_max_clicks=10): 11 | """“用户-物品”异构图 12 | 13 | https://docs.dgl.ai/en/latest/guide/training.html#heterogeneous-graphs 14 | 15 | 顶点类型 16 | ===== 17 | * user - 用户 18 | * item - 物品 19 | 20 | 边类型 21 | ===== 22 | * follow - 用户-关注->用户 23 | * follow-by - 用户-被关注->用户 24 | * click - 用户-点击->物品 25 | * click-by - 物品-被点击->用户 26 | * dislike - 用户-不喜欢->物品 27 | * dislike-by - 物品-被不喜欢->用户 28 | 29 | 顶点特征 30 | ===== 31 | user 32 | 33 | * feat - 维数为n_features 34 | * label - 范围[0, n_user_classes - 1] 35 | * train_mask 36 | 37 | item 38 | 39 | * feat - 维数为n_features 40 | 41 | 边特征 42 | ===== 43 | click 44 | 45 | * label - 范围[1, n_max_clicks] 46 | * train_mask 47 | 48 | :param n_users: “用户”顶点个数 49 | :param n_items: “物品”顶点个数 50 | :param n_follows: “用户-关注->用户”边条数 51 | :param n_clicks: “用户-点击->物品”边条数 52 | :param n_dislikes: “用户-不喜欢->物品”边条数 53 | :param n_features: “用户”和“物品”顶点特征维数 54 | :param n_user_classes: “用户”顶点标签类别数 55 | :param n_max_clicks: “用户-点击->物品”边标签类别数 56 | """ 57 | super().__init__('user-item') 58 | follow_src = np.random.randint(0, n_users, n_follows) 59 | follow_dst = np.random.randint(0, n_users, n_follows) 60 | click_src = np.random.randint(0, n_users, n_clicks) 61 | click_dst = np.random.randint(0, n_items, n_clicks) 62 | dislike_src = np.random.randint(0, n_users, n_dislikes) 63 | dislike_dst = np.random.randint(0, n_items, n_dislikes) 64 | 65 | g = dgl.heterograph({ 66 | ('user', 'follow', 'user'): (follow_src, follow_dst), 67 | ('user', 'followed-by', 'user'): (follow_dst, follow_src), 68 | ('user', 'click', 'item'): (click_src, click_dst), 69 | ('item', 'clicked-by', 'user'): (click_dst, click_src), 70 | ('user', 'dislike', 'item'): (dislike_src, dislike_dst), 71 | ('item', 'disliked-by', 'user'): (dislike_dst, dislike_src) 72 | }, num_nodes_dict={'user': n_users, 'item': n_items}) 73 | g.nodes['user'].data['feat'] = torch.randn(n_users, n_features) 74 | g.nodes['item'].data['feat'] = torch.randn(n_items, n_features) 75 | g.nodes['user'].data['label'] = torch.randint(0, n_user_classes, (n_users,)) 76 | g.edges['click'].data['label'] = torch.randint(1, n_max_clicks + 1, (n_clicks,)) 77 | # randomly generate training masks on user nodes and click edges 78 | g.nodes['user'].data['train_mask'] = torch.zeros(n_users, dtype=torch.bool).bernoulli(0.6) 79 | g.edges['click'].data['train_mask'] = torch.zeros(n_clicks, dtype=torch.bool).bernoulli(0.6) 80 | self.g = g 81 | 82 | def process(self): 83 | pass 84 | 85 | def __getitem__(self, idx): 86 | if idx != 0: 87 | raise IndexError('This dataset has only one graph') 88 | return self.g 89 | 90 | def __len__(self): 91 | return 1 92 | 93 | 94 | class RandomGraphDataset(DGLDataset): 95 | 96 | def __init__(self, num_nodes, num_edges, num_feats): 97 | """随机图数据集 98 | 99 | :param num_nodes: int 顶点数 100 | :param num_edges: int 边数 101 | :param num_feats: int 顶点特征维数 102 | """ 103 | super().__init__('random-graph') 104 | src = torch.randint(0, num_nodes, (num_edges,)) 105 | dst = torch.randint(0, num_nodes, (num_edges,)) 106 | # make it symmetric 107 | g = dgl.graph((torch.cat([src, dst]), torch.cat([dst, src])), num_nodes=num_nodes) 108 | # synthetic node features 109 | g.ndata['feat'] = torch.randn(num_nodes, num_feats) 110 | self.g = g 111 | 112 | def process(self): 113 | pass 114 | 115 | def __getitem__(self, idx): 116 | if idx != 0: 117 | raise IndexError('This dataset has only one graph') 118 | return self.g 119 | 120 | def __len__(self): 121 | return 1 122 | -------------------------------------------------------------------------------- /gnn/dgl/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/dgl/__init__.py -------------------------------------------------------------------------------- /gnn/dgl/dgl_first_demo.py: -------------------------------------------------------------------------------- 1 | """https://docs.dgl.ai/en/0.5.x/tutorials/basics/1_first.html""" 2 | import itertools 3 | 4 | import dgl 5 | import matplotlib.animation as animation 6 | import matplotlib.pyplot as plt 7 | import networkx as nx 8 | import numpy as np 9 | import torch 10 | import torch.nn as nn 11 | import torch.nn.functional as F 12 | 13 | from gnn.dgl.model import GCNFull 14 | 15 | 16 | def build_karate_club_graph(): 17 | src = np.array([ 18 | 1, 2, 2, 3, 3, 3, 4, 5, 6, 6, 6, 7, 7, 7, 7, 8, 8, 9, 10, 10, 19 | 10, 11, 12, 12, 13, 13, 13, 13, 16, 16, 17, 17, 19, 19, 21, 21, 25, 25, 27, 27, 20 | 27, 28, 29, 29, 30, 30, 31, 31, 31, 31, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 21 | 32, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33 22 | ]) 23 | dst = np.array([ 24 | 0, 0, 1, 0, 1, 2, 0, 0, 0, 4, 5, 0, 1, 2, 3, 0, 2, 2, 0, 4, 25 | 5, 0, 0, 3, 0, 1, 2, 3, 5, 6, 0, 1, 0, 1, 0, 1, 23, 24, 2, 23, 26 | 24, 2, 23, 26, 1, 8, 0, 24, 25, 28, 2, 8, 14, 15, 18, 20, 22, 23, 29, 30, 27 | 31, 8, 9, 13, 14, 15, 18, 19, 20, 22, 23, 26, 27, 28, 29, 30, 31, 32 28 | ]) 29 | return dgl.to_bidirected(dgl.graph((src, dst))) 30 | 31 | 32 | def draw_graph(g): 33 | nx_g = g.to_networkx().to_undirected() 34 | pos = nx.kamada_kawai_layout(nx_g) 35 | nx.draw(nx_g, pos, with_labels=True, node_color=[[.7, .7, .7]]) 36 | plt.show() 37 | 38 | 39 | def draw(i, g, ax, all_logits): 40 | cls1color = '#00FFFF' 41 | cls2color = '#FF00FF' 42 | pos = {} 43 | colors = [] 44 | for v in range(34): 45 | pos[v] = all_logits[i][v].numpy() 46 | cls = pos[v].argmax() 47 | colors.append(cls1color if cls else cls2color) 48 | ax.cla() 49 | ax.axis('off') 50 | ax.set_title('Epoch: %d' % i) 51 | nx.draw_networkx( 52 | g.to_networkx().to_undirected(), pos, node_color=colors, 53 | with_labels=True, node_size=300, ax=ax 54 | ) 55 | 56 | 57 | def main(): 58 | # Step 1: Creating a graph in DGL 59 | g = build_karate_club_graph() 60 | print('We have {} nodes.'.format(g.number_of_nodes())) 61 | print('We have {} edges.'.format(g.number_of_edges())) 62 | draw_graph(g) 63 | 64 | # Step 2: Assign features to nodes or edges 65 | embed = nn.Embedding(34, 5) # 34 nodes with embedding dim equal to 5 66 | g.ndata['feat'] = embed.weight 67 | print("node 2's input feature =", g.ndata['feat'][2]) 68 | 69 | # Step 3: Define a Graph Convolutional Network (GCN) 70 | # input feature size of 5 -> hidden size of 5 71 | # -> output feature size of 2 <=> 2 groups of the karate club 72 | net = GCNFull(5, 5, 2) 73 | 74 | # Step 4: Data preparation and initialization 75 | inputs = embed.weight 76 | labeled_nodes = torch.tensor([0, 33]) # only the instructor and the president nodes are labeled 77 | labels = torch.tensor([0, 1]) # their labels are different 78 | 79 | # Step 5: Train then visualize 80 | optimizer = torch.optim.Adam(itertools.chain(net.parameters(), embed.parameters()), lr=0.01) 81 | all_logits = [] 82 | for epoch in range(50): 83 | logits = net(g, inputs) 84 | # we save the logits for visualization later 85 | all_logits.append(logits.detach()) 86 | logp = F.log_softmax(logits, 1) 87 | # we only compute loss for labeled nodes 88 | loss = F.nll_loss(logp[labeled_nodes], labels) 89 | 90 | optimizer.zero_grad() 91 | loss.backward() 92 | optimizer.step() 93 | print('Epoch %d | Loss: %.4f' % (epoch, loss.item())) 94 | 95 | fig = plt.figure(dpi=150) 96 | fig.clf() 97 | ax = fig.subplots() 98 | # 在PyCharm中运行无法显示动画，删掉"ani="也无法显示。。 99 | ani = animation.FuncAnimation( 100 | fig, draw, frames=len(all_logits), fargs=(g, ax, all_logits), interval=200 101 | ) 102 | plt.show() 103 | 104 | 105 | if __name__ == '__main__': 106 | main() 107 | -------------------------------------------------------------------------------- /gnn/dgl/edge_clf.py: -------------------------------------------------------------------------------- 1 | """训练用于边分类/回归任务的GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/training-edge.html 4 | """ 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import torch.optim as optim 9 | 10 | from gnn.data import RandomGraphDataset 11 | from gnn.dgl.model import SAGEFull, DotProductPredictor 12 | 13 | 14 | class Model(nn.Module): 15 | 16 | def __init__(self, in_features, hidden_features, out_features): 17 | super().__init__() 18 | self.sage = SAGEFull(in_features, hidden_features, out_features) 19 | self.pred = DotProductPredictor() 20 | 21 | def forward(self, g, x): 22 | h = self.sage(g, x) 23 | return self.pred(g, h) 24 | 25 | 26 | def main(): 27 | data = RandomGraphDataset(100, 500, 10) 28 | g = data[0] 29 | # synthetic edge labels 30 | g.edata['label'] = torch.randn(1000) 31 | # synthetic train-validation-test splits 32 | g.edata['train_mask'] = torch.zeros(1000, dtype=torch.bool).bernoulli(0.6) 33 | 34 | edge_label = g.edata['label'] 35 | train_mask = g.edata['train_mask'] 36 | model = Model(10, 20, 5) 37 | opt = optim.Adam(model.parameters()) 38 | 39 | for epoch in range(30): 40 | pred = model(g, g.ndata['feat']) 41 | loss = F.mse_loss(pred[train_mask][:, 0], edge_label[train_mask]) 42 | opt.zero_grad() 43 | loss.backward() 44 | opt.step() 45 | print(loss.item()) 46 | 47 | 48 | if __name__ == '__main__': 49 | main() 50 | -------------------------------------------------------------------------------- /gnn/dgl/edge_clf_hetero.py: -------------------------------------------------------------------------------- 1 | """在异构图上训练用于边分类/回归任务的GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/training-edge.html 4 | """ 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | import torch.optim as optim 8 | 9 | from gnn.data import UserItemDataset 10 | from gnn.dgl.model import RGCNFull, HeteroDotProductPredictor 11 | 12 | 13 | class Model(nn.Module): 14 | 15 | def __init__(self, in_features, hidden_features, out_features, rel_names): 16 | super().__init__() 17 | self.rgcn = RGCNFull(in_features, hidden_features, out_features, rel_names) 18 | self.pred = HeteroDotProductPredictor() 19 | 20 | def forward(self, g, x, etype): 21 | h = self.rgcn(g, x) 22 | return self.pred(g, h, etype) 23 | 24 | 25 | def main(): 26 | data = UserItemDataset() 27 | g = data[0] 28 | in_feats = g.nodes['user'].data['feat'].shape[1] 29 | labels = g.edges['click'].data['label'].float() 30 | train_mask = g.edges['click'].data['train_mask'] 31 | 32 | model = Model(in_feats, 20, 5, g.etypes) 33 | opt = optim.Adam(model.parameters()) 34 | 35 | for epoch in range(10): 36 | pred = model(g, g.ndata['feat'], 'click') 37 | loss = F.mse_loss(pred[train_mask].squeeze(dim=1), labels[train_mask]) 38 | opt.zero_grad() 39 | loss.backward() 40 | opt.step() 41 | print(loss.item()) 42 | 43 | 44 | if __name__ == '__main__': 45 | main() 46 | -------------------------------------------------------------------------------- /gnn/dgl/edge_clf_hetero_mb.py: -------------------------------------------------------------------------------- 1 | """异构图上使用邻居采样的边分类GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/minibatch-edge.html 4 | """ 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | import torch.optim as optim 8 | from dgl.dataloading import MultiLayerFullNeighborSampler, EdgeDataLoader 9 | 10 | from gnn.data import UserItemDataset 11 | from gnn.dgl.model import RGCN, HeteroDotProductPredictor 12 | 13 | 14 | class Model(nn.Module): 15 | 16 | def __init__(self, in_features, hidden_features, out_features, rel_names): 17 | super().__init__() 18 | self.rgcn = RGCN(in_features, hidden_features, out_features, rel_names) 19 | self.pred = HeteroDotProductPredictor() 20 | 21 | def forward(self, edge_subgraph, blocks, x, etype): 22 | h = self.rgcn(blocks, x) 23 | return self.pred(edge_subgraph, h, etype) 24 | 25 | 26 | def main(): 27 | data = UserItemDataset() 28 | g = data[0] 29 | in_feats = g.nodes['user'].data['feat'].shape[1] 30 | g.edges['click'].data['label'] = g.edges['click'].data['label'].float() 31 | train_eids = g.edges['click'].data['train_mask'].nonzero(as_tuple=True)[0] 32 | 33 | sampler = MultiLayerFullNeighborSampler(2) 34 | dataloader = EdgeDataLoader(g, {'click': train_eids}, sampler, batch_size=256) 35 | 36 | model = Model(in_feats, 20, 5, g.etypes) 37 | optimizer = optim.Adam(model.parameters()) 38 | 39 | for epoch in range(10): 40 | model.train() 41 | losses = [] 42 | for input_nodes, edge_subgraph, blocks in dataloader: 43 | pred = model(edge_subgraph, blocks, blocks[0].srcdata['feat'], 'click') 44 | loss = F.mse_loss(pred.squeeze(dim=1), edge_subgraph.edges['click'].data['label']) 45 | losses.append(loss.item()) 46 | optimizer.zero_grad() 47 | loss.backward() 48 | optimizer.step() 49 | print('Epoch {:d} | Loss {:.4f}'.format(epoch, sum(losses) / len(losses))) 50 | 51 | 52 | if __name__ == '__main__': 53 | main() 54 | -------------------------------------------------------------------------------- /gnn/dgl/edge_clf_mb.py: -------------------------------------------------------------------------------- 1 | """使用邻居采样的边分类GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/minibatch-edge.html 4 | """ 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import torch.optim as optim 9 | from dgl.dataloading import MultiLayerFullNeighborSampler, EdgeDataLoader 10 | 11 | from gnn.data import RandomGraphDataset 12 | from gnn.dgl.model import GCN, MLPPredictor 13 | 14 | 15 | class Model(nn.Module): 16 | 17 | def __init__(self, in_features, hidden_features, out_features, num_classes): 18 | super().__init__() 19 | self.gcn = GCN(in_features, hidden_features, out_features) 20 | self.pred = MLPPredictor(out_features, num_classes) 21 | 22 | def forward(self, edge_subgraph, blocks, x): 23 | h = self.gcn(blocks, x) 24 | return self.pred(edge_subgraph, h) 25 | 26 | 27 | def main(): 28 | data = RandomGraphDataset(100, 500, 10) 29 | g = data[0] 30 | g.edata['label'] = torch.randint(0, 5, (g.num_edges(),)) 31 | train_mask = torch.zeros(g.num_edges(), dtype=torch.bool).bernoulli(0.6) 32 | train_idx = train_mask.nonzero(as_tuple=True)[0] 33 | 34 | sampler = MultiLayerFullNeighborSampler(2) 35 | dataloader = EdgeDataLoader(g, train_idx, sampler, batch_size=32) 36 | 37 | model = Model(10, 20, 10, 5) 38 | optimizer = optim.Adam(model.parameters()) 39 | 40 | for epoch in range(10): 41 | model.train() 42 | losses = [] 43 | for input_nodes, edge_subgraph, blocks in dataloader: 44 | logits = model(edge_subgraph, blocks, blocks[0].srcdata['feat']) 45 | loss = F.cross_entropy(logits, edge_subgraph.edata['label']) 46 | losses.append(loss.item()) 47 | optimizer.zero_grad() 48 | loss.backward() 49 | optimizer.step() 50 | print('Epoch {:d} | Loss {:.4f}'.format(epoch, sum(losses) / len(losses))) 51 | 52 | 53 | if __name__ == '__main__': 54 | main() 55 | -------------------------------------------------------------------------------- /gnn/dgl/edge_type_hetero.py: -------------------------------------------------------------------------------- 1 | """在异构图上训练用于预测边类型任务的GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/training-edge.html 4 | """ 5 | import dgl 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import torch.optim as optim 9 | 10 | from gnn.data import UserItemDataset 11 | from gnn.dgl.model import RGCNFull, MLPPredictor 12 | 13 | 14 | class Model(nn.Module): 15 | 16 | def __init__(self, in_features, hidden_features, out_features, out_classes, rel_names): 17 | super().__init__() 18 | self.rgcn = RGCNFull(in_features, hidden_features, out_features, rel_names) 19 | self.pred = MLPPredictor(out_features, out_classes) 20 | 21 | def forward(self, g, x, dec_graph): 22 | h = self.rgcn(g, x) 23 | return self.pred(dec_graph, h) 24 | 25 | 26 | def main(): 27 | data = UserItemDataset() 28 | g = data[0] 29 | in_feats = g.nodes['user'].data['feat'].shape[1] 30 | dec_graph = g['user', :, 'item'] 31 | edge_label = dec_graph.edata[dgl.ETYPE] 32 | edge_label -= edge_label.min().item() 33 | 34 | model = Model(in_feats, 20, 5, 2, g.etypes) 35 | opt = optim.Adam(model.parameters()) 36 | 37 | for epoch in range(10): 38 | logits = model(g, g.ndata['feat'], dec_graph) 39 | loss = F.cross_entropy(logits, edge_label) 40 | opt.zero_grad() 41 | loss.backward() 42 | opt.step() 43 | print(loss.item()) 44 | 45 | 46 | if __name__ == '__main__': 47 | main() 48 | -------------------------------------------------------------------------------- /gnn/dgl/graph_clf.py: -------------------------------------------------------------------------------- 1 | """训练用于图分类任务的GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/training-graph.html 4 | """ 5 | import dgl 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import torch.optim as optim 9 | from dgl.dataloading import GraphDataLoader 10 | 11 | from gnn.dgl.model import GCNFull 12 | 13 | 14 | class Model(nn.Module): 15 | 16 | def __init__(self, in_dim, hidden_dim, n_classes): 17 | super().__init__() 18 | self.gcn = GCNFull(in_dim, hidden_dim, hidden_dim) 19 | self.classify = nn.Linear(hidden_dim, n_classes) 20 | 21 | def forward(self, g, h): 22 | # Apply graph convolution and activation. 23 | h = self.gcn(g, h) 24 | with g.local_scope(): 25 | g.ndata['h'] = h 26 | # Calculate graph representation by average readout. 27 | hg = dgl.mean_nodes(g, 'h') 28 | return self.classify(hg) 29 | 30 | 31 | def main(): 32 | dataset = dgl.data.GINDataset('MUTAG', False) 33 | dataloader = GraphDataLoader(dataset, batch_size=32, shuffle=True) 34 | 35 | model = Model(dataset.dim_nfeats, 20, dataset.gclasses) 36 | opt = optim.Adam(model.parameters()) 37 | 38 | for epoch in range(5): 39 | for batched_graph, labels in dataloader: 40 | feats = batched_graph.ndata['attr'].float() 41 | logits = model(batched_graph, feats) 42 | loss = F.cross_entropy(logits, labels) 43 | opt.zero_grad() 44 | loss.backward() 45 | opt.step() 46 | print(loss.item()) 47 | 48 | 49 | if __name__ == '__main__': 50 | main() 51 | -------------------------------------------------------------------------------- /gnn/dgl/graph_clf_hetero.py: -------------------------------------------------------------------------------- 1 | """在异构图上训练用于图分类任务的GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/training-graph.html 4 | """ 5 | import random 6 | 7 | import dgl 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | import torch.optim as optim 11 | from dgl.data import DGLDataset 12 | from dgl.dataloading import GraphDataLoader 13 | 14 | from gnn.data import UserItemDataset as UserItemOneGraphDataset 15 | from gnn.dgl.model import RGCNFull 16 | 17 | 18 | class HeteroClassifier(nn.Module): 19 | 20 | def __init__(self, in_dim, hidden_dim, n_classes, rel_names): 21 | super().__init__() 22 | self.rgcn = RGCNFull(in_dim, hidden_dim, hidden_dim, rel_names) 23 | self.classify = nn.Linear(hidden_dim, n_classes) 24 | 25 | def forward(self, g): 26 | h = self.rgcn(g, g.ndata['feat']) 27 | with g.local_scope(): 28 | g.ndata['h'] = h 29 | # Calculate graph representation by average readout. 30 | hg = sum(dgl.mean_nodes(g, 'h', ntype=ntype) for ntype in g.ntypes) 31 | return F.softmax(self.classify(hg), dim=0) 32 | 33 | 34 | class UserItemDataset(DGLDataset): 35 | 36 | def __init__(self): 37 | self.graphs = [] 38 | self.labels = [] 39 | self.n_features = 10 40 | self.num_classes = 3 41 | super().__init__('user-item') 42 | 43 | def process(self): 44 | random.seed(42) 45 | # 随机构造，无实际意义 46 | for i in range(100): 47 | g = UserItemOneGraphDataset( 48 | n_users=random.randint(5, 10), 49 | n_items=random.randint(5, 10), 50 | n_follows=random.randint(10, 20), 51 | n_clicks=random.randint(10, 20), 52 | n_dislikes=random.randint(10, 20), 53 | n_features=self.n_features 54 | )[0] 55 | self.graphs.append(g) 56 | self.labels.append(random.randrange(self.num_classes)) 57 | 58 | def __getitem__(self, idx): 59 | return self.graphs[idx], self.labels[idx] 60 | 61 | def __len__(self): 62 | return len(self.graphs) 63 | 64 | 65 | def main(): 66 | dataset = UserItemDataset() 67 | dataloader = GraphDataLoader(dataset, batch_size=32, shuffle=True) 68 | 69 | model = HeteroClassifier(dataset.n_features, 20, dataset.num_classes, dataset[0][0].etypes) 70 | opt = optim.Adam(model.parameters()) 71 | 72 | for epoch in range(5): 73 | for batched_graph, labels in dataloader: 74 | logits = model(batched_graph) 75 | loss = F.cross_entropy(logits, labels) 76 | opt.zero_grad() 77 | loss.backward() 78 | opt.step() 79 | print(loss.item()) 80 | 81 | 82 | if __name__ == '__main__': 83 | main() 84 | -------------------------------------------------------------------------------- /gnn/dgl/link_pred.py: -------------------------------------------------------------------------------- 1 | """训练用于连接预测任务的GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/training-link.html 4 | """ 5 | import dgl 6 | import torch 7 | import torch.nn as nn 8 | import torch.optim as optim 9 | 10 | from gnn.data import RandomGraphDataset 11 | from gnn.dgl.model import SAGEFull, DotProductPredictor, MarginLoss 12 | 13 | 14 | def construct_negative_graph(graph, k): 15 | src, dst = graph.edges() 16 | neg_src = src.repeat_interleave(k) 17 | neg_dst = torch.randint(0, graph.number_of_nodes(), (len(src) * k,)) 18 | return dgl.graph((neg_src, neg_dst), num_nodes=graph.number_of_nodes()) 19 | 20 | 21 | class Model(nn.Module): 22 | 23 | def __init__(self, in_features, hidden_features, out_features): 24 | super().__init__() 25 | self.sage = SAGEFull(in_features, hidden_features, out_features) 26 | self.pred = DotProductPredictor() 27 | 28 | def forward(self, g, neg_g, x): 29 | h = self.sage(g, x) 30 | return self.pred(g, h), self.pred(neg_g, h) 31 | 32 | 33 | def main(): 34 | data = RandomGraphDataset(100, 500, 10) 35 | g = data[0] 36 | node_features = g.ndata['feat'] 37 | n_features = node_features.shape[1] 38 | k = 5 39 | 40 | model = Model(n_features, 100, 100) 41 | opt = optim.Adam(model.parameters()) 42 | loss_func = MarginLoss() 43 | 44 | for epoch in range(10): 45 | negative_graph = construct_negative_graph(g, k) 46 | pos_score, neg_score = model(g, negative_graph, node_features) 47 | loss = loss_func(pos_score, neg_score) 48 | opt.zero_grad() 49 | loss.backward() 50 | opt.step() 51 | print(loss.item()) 52 | 53 | 54 | if __name__ == '__main__': 55 | main() 56 | -------------------------------------------------------------------------------- /gnn/dgl/link_pred_hetero.py: -------------------------------------------------------------------------------- 1 | """在异构图上训练用于连接预测任务的GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/training-link.html 4 | """ 5 | import dgl 6 | import torch 7 | import torch.nn as nn 8 | import torch.optim as optim 9 | 10 | from gnn.data import UserItemDataset 11 | from gnn.dgl.model import RGCNFull, HeteroDotProductPredictor, MarginLoss 12 | 13 | 14 | def construct_negative_graph(graph, k, etype): 15 | utype, _, vtype = etype 16 | src, dst = graph.edges(etype=etype) 17 | neg_src = src.repeat_interleave(k) 18 | neg_dst = torch.randint(0, graph.number_of_nodes(vtype), (len(src) * k,)) 19 | return dgl.heterograph( 20 | {etype: (neg_src, neg_dst)}, 21 | num_nodes_dict={ntype: graph.number_of_nodes(ntype) for ntype in graph.ntypes} 22 | ) 23 | 24 | 25 | class Model(nn.Module): 26 | 27 | def __init__(self, in_features, hidden_features, out_features, rel_names): 28 | super().__init__() 29 | self.rgcn = RGCNFull(in_features, hidden_features, out_features, rel_names) 30 | self.pred = HeteroDotProductPredictor() 31 | 32 | def forward(self, g, neg_g, x, etype): 33 | h = self.rgcn(g, x) 34 | return self.pred(g, h, etype), self.pred(neg_g, h, etype) 35 | 36 | 37 | def main(): 38 | data = UserItemDataset() 39 | g = data[0] 40 | in_feats = g.nodes['user'].data['feat'].shape[1] 41 | k = 5 42 | 43 | model = Model(in_feats, 20, 5, g.etypes) 44 | opt = optim.Adam(model.parameters()) 45 | loss_func = MarginLoss() 46 | 47 | for epoch in range(10): 48 | negative_graph = construct_negative_graph(g, k, ('user', 'click', 'item')) 49 | pos_score, neg_score = model( 50 | g, negative_graph, g.ndata['feat'], ('user', 'click', 'item') 51 | ) 52 | loss = loss_func(pos_score, neg_score) 53 | opt.zero_grad() 54 | loss.backward() 55 | opt.step() 56 | print(loss.item()) 57 | 58 | 59 | if __name__ == '__main__': 60 | main() 61 | -------------------------------------------------------------------------------- /gnn/dgl/link_pred_hetero_mb.py: -------------------------------------------------------------------------------- 1 | """异构图上使用邻居采样的连接预测GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/minibatch-link.html 4 | """ 5 | import torch.nn as nn 6 | import torch.optim as optim 7 | from dgl.dataloading import MultiLayerFullNeighborSampler, EdgeDataLoader 8 | from dgl.dataloading.negative_sampler import Uniform 9 | 10 | from gnn.data import UserItemDataset 11 | from gnn.dgl.model import RGCN, HeteroDotProductPredictor, MarginLoss 12 | 13 | 14 | class Model(nn.Module): 15 | 16 | def __init__(self, in_features, hidden_features, out_features, rel_names): 17 | super().__init__() 18 | self.rgcn = RGCN(in_features, hidden_features, out_features, rel_names) 19 | self.pred = HeteroDotProductPredictor() 20 | 21 | def forward(self, pos_g, neg_g, blocks, x, etype): 22 | h = self.rgcn(blocks, x) 23 | return self.pred(pos_g, h, etype), self.pred(neg_g, h, etype) 24 | 25 | 26 | def main(): 27 | data = UserItemDataset() 28 | g = data[0] 29 | in_feats = g.nodes['user'].data['feat'].shape[1] 30 | train_eids = g.edges['click'].data['train_mask'].nonzero(as_tuple=True)[0] 31 | 32 | sampler = MultiLayerFullNeighborSampler(2) 33 | dataloader = EdgeDataLoader( 34 | g, {'click': train_eids}, sampler, negative_sampler=Uniform(5), batch_size=256 35 | ) 36 | 37 | model = Model(in_feats, 20, 5, g.etypes) 38 | optimizer = optim.Adam(model.parameters()) 39 | loss_func = MarginLoss() 40 | 41 | for epoch in range(10): 42 | model.train() 43 | losses = [] 44 | for input_nodes, pos_g, neg_g, blocks in dataloader: 45 | pos_score, neg_score = model(pos_g, neg_g, blocks, blocks[0].srcdata['feat'], 'click') 46 | loss = loss_func(pos_score, neg_score) 47 | losses.append(loss.item()) 48 | optimizer.zero_grad() 49 | loss.backward() 50 | optimizer.step() 51 | print('Epoch {:d} | Loss {:.4f}'.format(epoch, sum(losses) / len(losses))) 52 | 53 | 54 | if __name__ == '__main__': 55 | main() 56 | -------------------------------------------------------------------------------- /gnn/dgl/link_pred_mb.py: -------------------------------------------------------------------------------- 1 | """使用邻居采样的连接预测GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/minibatch-link.html 4 | """ 5 | import torch 6 | import torch.nn as nn 7 | import torch.optim as optim 8 | from dgl.dataloading import MultiLayerFullNeighborSampler, EdgeDataLoader 9 | from dgl.dataloading.negative_sampler import Uniform 10 | 11 | from gnn.data import RandomGraphDataset 12 | from gnn.dgl.model import GCN, DotProductPredictor, MarginLoss 13 | 14 | 15 | class Model(nn.Module): 16 | 17 | def __init__(self, in_features, hidden_features, out_features): 18 | super().__init__() 19 | self.gcn = GCN(in_features, hidden_features, out_features) 20 | self.pred = DotProductPredictor() 21 | 22 | def forward(self, pos_g, neg_g, blocks, x): 23 | h = self.gcn(blocks, x) 24 | return self.pred(pos_g, h), self.pred(neg_g, h) 25 | 26 | 27 | def main(): 28 | data = RandomGraphDataset(100, 500, 10) 29 | g = data[0] 30 | train_mask = torch.zeros(g.num_edges(), dtype=torch.bool).bernoulli(0.6) 31 | train_idx = train_mask.nonzero(as_tuple=True)[0] 32 | 33 | sampler = MultiLayerFullNeighborSampler(2) 34 | dataloader = EdgeDataLoader(g, train_idx, sampler, negative_sampler=Uniform(5), batch_size=32) 35 | 36 | model = Model(10, 100, 10) 37 | optimizer = optim.Adam(model.parameters()) 38 | loss_func = MarginLoss() 39 | 40 | for epoch in range(10): 41 | model.train() 42 | losses = [] 43 | for input_nodes, pos_g, neg_g, blocks in dataloader: 44 | pos_score, neg_score = model(pos_g, neg_g, blocks, blocks[0].srcdata['feat']) 45 | loss = loss_func(pos_score, neg_score) 46 | losses.append(loss.item()) 47 | optimizer.zero_grad() 48 | loss.backward() 49 | optimizer.step() 50 | print('Epoch {:d} | Loss {:.4f}'.format(epoch, sum(losses) / len(losses))) 51 | 52 | 53 | if __name__ == '__main__': 54 | main() 55 | -------------------------------------------------------------------------------- /gnn/dgl/model.py: -------------------------------------------------------------------------------- 1 | import dgl.function as fn 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from dgl.nn import GraphConv, SAGEConv, HeteroGraphConv 6 | 7 | 8 | class GCNFull(nn.Module): 9 | 10 | def __init__(self, in_feats, hidden_size, num_classes): 11 | super().__init__() 12 | self.conv1 = GraphConv(in_feats, hidden_size) 13 | self.conv2 = GraphConv(hidden_size, num_classes) 14 | 15 | def forward(self, g, inputs): 16 | h = self.conv1(g, inputs) 17 | h = F.relu(h) 18 | h = self.conv2(g, h) 19 | return h 20 | 21 | 22 | class GCN(nn.Module): 23 | 24 | def __init__(self, in_feats, hidden_size, num_classes): 25 | super().__init__() 26 | self.conv1 = GraphConv(in_feats, hidden_size) 27 | self.conv2 = GraphConv(hidden_size, num_classes) 28 | 29 | def forward(self, blocks, inputs): 30 | h = self.conv1(blocks[0], inputs) 31 | h = F.relu(h) 32 | h = self.conv2(blocks[1], h) 33 | return h 34 | 35 | 36 | class SAGEFull(nn.Module): 37 | 38 | def __init__(self, in_feats, hid_feats, out_feats): 39 | super().__init__() 40 | self.conv1 = SAGEConv(in_feats, hid_feats, 'mean') 41 | self.conv2 = SAGEConv(hid_feats, out_feats, 'mean') 42 | 43 | def forward(self, g, inputs): 44 | # inputs are features of nodes 45 | h = F.relu(self.conv1(g, inputs)) 46 | h = self.conv2(g, h) 47 | return h 48 | 49 | 50 | class RGCNFull(nn.Module): 51 | 52 | def __init__(self, in_feats, hid_feats, out_feats, rel_names): 53 | super().__init__() 54 | self.conv1 = HeteroGraphConv({ 55 | rel: GraphConv(in_feats, hid_feats) for rel in rel_names 56 | }, aggregate='sum') 57 | self.conv2 = HeteroGraphConv({ 58 | rel: GraphConv(hid_feats, out_feats) for rel in rel_names 59 | }, aggregate='sum') 60 | 61 | def forward(self, g, inputs): 62 | h = self.conv1(g, inputs) 63 | h = {k: F.relu(v) for k, v in h.items()} 64 | h = self.conv2(g, h) 65 | return h 66 | 67 | 68 | class RGCN(nn.Module): 69 | 70 | def __init__(self, in_feats, hid_feats, out_feats, rel_names): 71 | super().__init__() 72 | self.conv1 = HeteroGraphConv({ 73 | rel: GraphConv(in_feats, hid_feats) for rel in rel_names 74 | }, aggregate='sum') 75 | self.conv2 = HeteroGraphConv({ 76 | rel: GraphConv(hid_feats, out_feats) for rel in rel_names 77 | }, aggregate='sum') 78 | 79 | def forward(self, blocks, inputs): 80 | h = self.conv1(blocks[0], inputs) 81 | h = {k: F.relu(v) for k, v in h.items()} 82 | h = self.conv2(blocks[1], h) 83 | return h 84 | 85 | 86 | class DotProductPredictor(nn.Module): 87 | 88 | def forward(self, graph, h): 89 | # h contains the node representations computed from the GNN defined in node_clf.py 90 | with graph.local_scope(): 91 | graph.ndata['h'] = h 92 | graph.apply_edges(fn.u_dot_v('h', 'h', 'score')) 93 | return graph.edata['score'] 94 | 95 | 96 | class MLPPredictor(nn.Module): 97 | 98 | def __init__(self, in_features, out_classes): 99 | super().__init__() 100 | self.W = nn.Linear(in_features * 2, out_classes) 101 | 102 | def apply_edges(self, edges): 103 | score = self.W(torch.cat([edges.src['h'], edges.dst['h']], dim=1)) 104 | return {'score': score} 105 | 106 | def forward(self, graph, h): 107 | # h contains the node representations computed from the GNN defined in node_clf.py 108 | with graph.local_scope(): 109 | graph.ndata['h'] = h 110 | graph.apply_edges(self.apply_edges) 111 | return graph.edata['score'] 112 | 113 | 114 | class HeteroDotProductPredictor(nn.Module): 115 | 116 | def forward(self, graph, h, etype): 117 | # h contains the node representations for each edge type computed from node_clf_hetero.py 118 | with graph.local_scope(): 119 | graph.ndata['h'] = h # assigns 'h' of all node types in one shot 120 | graph.apply_edges(fn.u_dot_v('h', 'h', 'score'), etype=etype) 121 | return graph.edges[etype].data['score'] 122 | 123 | 124 | class HeteroMLPPredictor(nn.Module): 125 | 126 | def __init__(self, in_features, out_classes): 127 | super().__init__() 128 | self.W = nn.Linear(in_features * 2, out_classes) 129 | 130 | def apply_edges(self, edges): 131 | score = self.W(torch.cat([edges.src['h'], edges.dst['h']], dim=1)) 132 | return {'score': score} 133 | 134 | def forward(self, graph, h, etype): 135 | # h contains the node representations for each edge type computed from node_clf_hetero.py 136 | with graph.local_scope(): 137 | graph.ndata['h'] = h # assigns 'h' of all node types in one shot 138 | graph.apply_edges(self.apply_edges, etype=etype) 139 | return graph.edges[etype].data['score'] 140 | 141 | 142 | class MarginLoss(nn.Module): 143 | 144 | def forward(self, pos_score, neg_score): 145 | return (1 - pos_score + neg_score.view(pos_score.shape[0], -1)).clamp(min=0).mean() 146 | -------------------------------------------------------------------------------- /gnn/dgl/node_clf.py: -------------------------------------------------------------------------------- 1 | """训练用于顶点分类/回归任务的GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/training-node.html 4 | """ 5 | import torch.nn.functional as F 6 | import torch.optim as optim 7 | from dgl.data import CiteseerGraphDataset 8 | 9 | from gnn.dgl.model import SAGEFull 10 | from gnn.utils import accuracy 11 | 12 | 13 | def main(): 14 | # load data 15 | dataset = CiteseerGraphDataset() 16 | g = dataset[0] 17 | 18 | node_features = g.ndata['feat'] 19 | node_labels = g.ndata['label'] 20 | n_features = node_features.shape[1] 21 | n_labels = dataset.num_classes 22 | 23 | # get split masks 24 | train_mask = g.ndata['train_mask'] 25 | valid_mask = g.ndata['val_mask'] 26 | test_mask = g.ndata['test_mask'] 27 | 28 | model = SAGEFull(in_feats=n_features, hid_feats=100, out_feats=n_labels) 29 | opt = optim.Adam(model.parameters()) 30 | 31 | for epoch in range(30): 32 | model.train() 33 | # forward propagation by using all nodes 34 | logits = model(g, node_features) 35 | # compute loss 36 | loss = F.cross_entropy(logits[train_mask], node_labels[train_mask]) 37 | # compute validation accuracy 38 | acc = accuracy(logits[valid_mask], node_labels[valid_mask]) 39 | # backward propagation 40 | opt.zero_grad() 41 | loss.backward() 42 | opt.step() 43 | print('Epoch {:d} | Loss {:.4f} | Accuracy {:.2%}'.format(epoch + 1, loss.item(), acc)) 44 | acc = accuracy(model(g, node_features)[test_mask], node_labels[test_mask]) 45 | print('Test accuracy {:.4f}'.format(acc)) 46 | 47 | 48 | if __name__ == '__main__': 49 | main() 50 | -------------------------------------------------------------------------------- /gnn/dgl/node_clf_hetero.py: -------------------------------------------------------------------------------- 1 | """在异构图上训练用于顶点分类/回归任务的GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/training-node.html 4 | """ 5 | import torch.nn.functional as F 6 | import torch.optim as optim 7 | 8 | from gnn.data import UserItemDataset 9 | from gnn.dgl.model import RGCNFull 10 | 11 | 12 | def main(): 13 | data = UserItemDataset() 14 | g = data[0] 15 | in_feats = g.nodes['user'].data['feat'].shape[1] 16 | labels = g.nodes['user'].data['label'] 17 | train_mask = g.nodes['user'].data['train_mask'] 18 | 19 | model = RGCNFull(in_feats, 20, labels.max().item() + 1, g.etypes) 20 | opt = optim.Adam(model.parameters()) 21 | 22 | for epoch in range(10): 23 | model.train() 24 | # forward propagation by using all nodes and extracting the user embeddings 25 | logits = model(g, g.ndata['feat'])['user'] 26 | # compute loss 27 | loss = F.cross_entropy(logits[train_mask], labels[train_mask]) 28 | # compute validation accuracy, omitted in this example 29 | # backward propagation 30 | opt.zero_grad() 31 | loss.backward() 32 | opt.step() 33 | print(loss.item()) 34 | 35 | 36 | if __name__ == '__main__': 37 | main() 38 | -------------------------------------------------------------------------------- /gnn/dgl/node_clf_hetero_mb.py: -------------------------------------------------------------------------------- 1 | """异构图上使用邻居采样的顶点分类GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/minibatch-node.html 4 | """ 5 | import torch 6 | import torch.nn.functional as F 7 | import torch.optim as optim 8 | from dgl.dataloading import MultiLayerFullNeighborSampler, NodeDataLoader 9 | 10 | from gnn.data import UserItemDataset 11 | from gnn.dgl.model import RGCN 12 | 13 | 14 | def main(): 15 | data = UserItemDataset() 16 | g = data[0] 17 | train_idx = g.nodes['user'].data['train_mask'].nonzero(as_tuple=True)[0] 18 | 19 | sampler = MultiLayerFullNeighborSampler(2) 20 | dataloader = NodeDataLoader(g, {'user': train_idx}, sampler, batch_size=256) 21 | 22 | model = RGCN(10, 20, 5, g.etypes) 23 | opt = optim.Adam(model.parameters()) 24 | 25 | for epoch in range(10): 26 | model.train() 27 | losses = [] 28 | for input_nodes, output_nodes, blocks in dataloader: 29 | logits = model(blocks, blocks[0].srcdata['feat'])['user'] 30 | loss = F.cross_entropy(logits, blocks[-1].dstnodes['user'].data['label']) 31 | losses.append(loss.item()) 32 | opt.zero_grad() 33 | loss.backward() 34 | opt.step() 35 | print('Epoch {:d} | Loss {:.4f}'.format(epoch + 1, torch.tensor(losses).mean().item())) 36 | 37 | 38 | if __name__ == '__main__': 39 | main() 40 | -------------------------------------------------------------------------------- /gnn/dgl/node_clf_mb.py: -------------------------------------------------------------------------------- 1 | """使用邻居采样的顶点分类GNN 2 | 3 | https://docs.dgl.ai/en/latest/guide/minibatch-node.html 4 | """ 5 | import torch 6 | import torch.nn.functional as F 7 | import torch.optim as optim 8 | from dgl.data import CiteseerGraphDataset 9 | from dgl.dataloading import MultiLayerFullNeighborSampler, NodeDataLoader 10 | 11 | from gnn.dgl.model import GCN 12 | 13 | 14 | def main(): 15 | data = CiteseerGraphDataset() 16 | g = data[0] 17 | train_idx = g.ndata['train_mask'].nonzero(as_tuple=True)[0] 18 | 19 | sampler = MultiLayerFullNeighborSampler(2) 20 | dataloader = NodeDataLoader(g, train_idx, sampler, batch_size=32) 21 | 22 | model = GCN(g.ndata['feat'].shape[1], 100, data.num_classes) 23 | optimizer = optim.Adam(model.parameters()) 24 | for epoch in range(30): 25 | model.train() 26 | losses = [] 27 | for input_nodes, output_nodes, blocks in dataloader: 28 | logits = model(blocks, blocks[0].srcdata['feat']) 29 | loss = F.cross_entropy(logits, blocks[-1].dstdata['label']) 30 | losses.append(loss.item()) 31 | optimizer.zero_grad() 32 | loss.backward() 33 | optimizer.step() 34 | print('Epoch {:d} | Loss {:.4f}'.format(epoch + 1, torch.tensor(losses).mean().item())) 35 | 36 | 37 | if __name__ == '__main__': 38 | main() 39 | -------------------------------------------------------------------------------- /gnn/gat/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/gat/__init__.py -------------------------------------------------------------------------------- /gnn/gat/model.py: -------------------------------------------------------------------------------- 1 | """Graph Attention Networks (GAT) 2 | 3 | 论文链接：https://arxiv.org/abs/1710.10903 4 | """ 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from dgl.nn import GATConv 8 | 9 | 10 | class GAT(nn.Module): 11 | 12 | def __init__( 13 | self, in_dim, hidden_dim, out_dim, num_heads, dropout, 14 | residual=False, activation=None): 15 | """GAT模型 16 | 17 | :param in_dim: int 输入特征维数 18 | :param hidden_dim: int 隐含特征维数 19 | :param out_dim: int 输出特征维数 20 | :param num_heads: List[int] 每一层的注意力头数，长度等于层数 21 | :param dropout: float Dropout概率 22 | :param residual: bool, optional 是否使用残差连接，默认为False 23 | :param activation: callable, optional 输出层激活函数 24 | :raise ValueError: 如果层数（即num_heads的长度）小于2 25 | """ 26 | super().__init__() 27 | num_layers = len(num_heads) 28 | if num_layers < 2: 29 | raise ValueError('层数至少为2，实际为{}'.format(num_layers)) 30 | self.layers = nn.ModuleList() 31 | self.layers.append(GATConv( 32 | in_dim, hidden_dim, num_heads[0], dropout, dropout, residual=residual, activation=F.elu 33 | )) 34 | for i in range(1, num_layers - 1): 35 | self.layers.append(GATConv( 36 | num_heads[i - 1] * hidden_dim, hidden_dim, num_heads[i], dropout, dropout, 37 | residual=residual, activation=F.elu 38 | )) 39 | self.layers.append(GATConv( 40 | num_heads[-2] * hidden_dim, out_dim, num_heads[-1], dropout, dropout, 41 | residual=residual, activation=activation 42 | )) 43 | 44 | def forward(self, g, h): 45 | """ 46 | :param g: DGLGraph 同构图 47 | :param h: tensor(N, d_in) 输入特征，N为g的顶点数 48 | :return: tensor(N, d_out) 输出顶点特征，K为注意力头数 49 | """ 50 | for i in range(len(self.layers) - 1): 51 | h = self.layers[i](g, h).flatten(start_dim=1) # (N, K, d_hid) -> (N, K*d_hid) 52 | h = self.layers[-1](g, h).mean(dim=1) # (N, K, d_out) -> (N, d_out) 53 | return h 54 | -------------------------------------------------------------------------------- /gnn/gat/readme.txt: -------------------------------------------------------------------------------- 1 | 直推式 2 | Cora 3 | python -m gnn.gat.train_transductive --dataset=cora 4 | Test Accuracy 0.8230 5 | 6 | Citeseer 7 | python -m gnn.gat.train_transductive --dataset=citeseer 8 | Test Accuracy 0.7010 9 | 10 | Pubmed 11 | python -m gnn.gat.train_transductive --dataset=pubmed --num-out-heads=8 --lr=0.01 --weight-decay=0.001 12 | Test Accuracy 0.7950 13 | 14 | 归纳式 15 | PPI 16 | python -m gnn.gat.train_inductive 17 | Test F1-score 0.9863 18 | -------------------------------------------------------------------------------- /gnn/gat/train_inductive.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import dgl 4 | import numpy as np 5 | import torch 6 | import torch.nn.functional as F 7 | import torch.optim as optim 8 | from dgl.data import PPIDataset 9 | from sklearn.metrics import f1_score 10 | from torch.utils.data.dataloader import DataLoader 11 | 12 | from gnn.gat.model import GAT 13 | 14 | 15 | def train(args): 16 | train_dataset = PPIDataset('train') 17 | valid_dataset = PPIDataset('valid') 18 | test_dataset = PPIDataset('test') 19 | train_dataloader = DataLoader(train_dataset, args.batch_size, collate_fn=dgl.batch) 20 | valid_dataloader = DataLoader(valid_dataset, args.batch_size, collate_fn=dgl.batch) 21 | test_dataloader = DataLoader(test_dataset, args.batch_size, collate_fn=dgl.batch) 22 | 23 | num_feats = train_dataset[0].ndata['feat'].shape[1] 24 | num_classes = train_dataset.num_labels 25 | 26 | num_heads = [args.num_heads] * (args.num_layers - 1) + [args.num_out_heads] 27 | model = GAT(num_feats, args.num_hidden, num_classes, num_heads, args.dropout, residual=True) 28 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 29 | for epoch in range(args.epochs): 30 | model.train() 31 | losses = [] 32 | train_scores = [] 33 | for bg in train_dataloader: 34 | logits = model(bg, bg.ndata['feat']) 35 | labels = bg.ndata['label'].float() 36 | loss = F.binary_cross_entropy_with_logits(logits, labels) 37 | losses.append(loss.item()) 38 | train_scores.append(micro_f1_score(logits, labels)) 39 | optimizer.zero_grad() 40 | loss.backward() 41 | optimizer.step() 42 | print('Epoch {:04d} | Loss {:.4f} | Train F1-score {:.4f}'.format( 43 | epoch, np.array(losses).mean(), np.array(train_scores).mean() 44 | )) 45 | if epoch % 5 == 0: 46 | val_scores = [ 47 | evaluate(model, bg, bg.ndata['feat'], bg.ndata['label']) for bg in valid_dataloader 48 | ] 49 | print('Val F1-score {:.4f}'.format(np.array(val_scores).mean())) 50 | 51 | print() 52 | test_scores = [ 53 | evaluate(model, bg, bg.ndata['feat'], bg.ndata['label']) for bg in test_dataloader 54 | ] 55 | print('Test F1-score {:.4f}'.format(np.array(test_scores).mean())) 56 | 57 | 58 | def micro_f1_score(logits, labels): 59 | predict = np.where(logits.detach().numpy() >= 0., 1, 0) 60 | return f1_score(labels.numpy(), predict, average='micro') 61 | 62 | 63 | def evaluate(model, g, feats, labels): 64 | with torch.no_grad(): 65 | model.eval() 66 | logits = model(g, feats) 67 | return micro_f1_score(logits, labels) 68 | 69 | 70 | def main(): 71 | parser = argparse.ArgumentParser(description='GAT Inductive Training') 72 | parser.add_argument('--batch-size', type=int, default=2, help='batch size') 73 | parser.add_argument('--num-layers', type=int, default=3, help='number of GAT layers') 74 | parser.add_argument('--num-hidden', type=int, default=256, help='number of hidden units') 75 | parser.add_argument( 76 | '--num-heads', type=int, default=4, help='number of attention heads in hidden layers' 77 | ) 78 | parser.add_argument( 79 | '--num-out-heads', type=int, default=6, help='number of attention heads in output layer' 80 | ) 81 | parser.add_argument('--dropout', type=float, default=0., help='dropout probability') 82 | parser.add_argument('--epochs', type=int, default=50, help='number of training epochs') 83 | parser.add_argument('--lr', type=float, default=0.005, help='learning rate') 84 | parser.add_argument('--weight-decay', type=float, default=0., help='weight decay') 85 | args = parser.parse_args() 86 | train(args) 87 | 88 | 89 | if __name__ == '__main__': 90 | main() 91 | -------------------------------------------------------------------------------- /gnn/gat/train_transductive.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import dgl 4 | import torch 5 | import torch.nn.functional as F 6 | import torch.optim as optim 7 | 8 | from gnn.gat.model import GAT 9 | from gnn.utils import load_citation_dataset, accuracy 10 | 11 | 12 | def train(args): 13 | data = load_citation_dataset(args.dataset) 14 | g = data[0] 15 | features = g.ndata['feat'] 16 | labels = g.ndata['label'] 17 | train_mask = g.ndata['train_mask'] 18 | val_mask = g.ndata['val_mask'] 19 | test_mask = g.ndata['test_mask'] 20 | 21 | # add self loop 22 | g = dgl.remove_self_loop(g) 23 | g = dgl.add_self_loop(g) 24 | 25 | num_heads = [args.num_heads] * (args.num_layers - 1) + [args.num_out_heads] 26 | model = GAT(features.shape[1], args.num_hidden, data.num_classes, num_heads, args.dropout) 27 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 28 | for epoch in range(args.epochs): 29 | model.train() 30 | logits = model(g, features) 31 | loss = F.cross_entropy(logits[train_mask], labels[train_mask]) 32 | optimizer.zero_grad() 33 | loss.backward() 34 | optimizer.step() 35 | 36 | train_acc = accuracy(logits[train_mask], labels[train_mask]) 37 | val_acc = evaluate(model, g, features, labels, val_mask) 38 | print('Epoch {:04d} | Loss {:.4f} | Train Acc {:.4f} | Val Acc {:.4f}'.format( 39 | epoch, loss.item(), train_acc, val_acc 40 | )) 41 | 42 | print() 43 | acc = evaluate(model, g, features, labels, test_mask) 44 | print('Test Accuracy {:.4f}'.format(acc)) 45 | 46 | 47 | def evaluate(model, g, features, labels, mask): 48 | model.eval() 49 | with torch.no_grad(): 50 | logits = model(g, features) 51 | return accuracy(logits[mask], labels[mask]) 52 | 53 | 54 | def main(): 55 | parser = argparse.ArgumentParser(description='GAT Transductive Training') 56 | parser.add_argument( 57 | '--dataset', choices=['cora', 'citeseer', 'pubmed'], default='cora', help='dataset' 58 | ) 59 | parser.add_argument('--num-layers', type=int, default=2, help='number of GAT layers') 60 | parser.add_argument('--num-hidden', type=int, default=8, help='number of hidden units') 61 | parser.add_argument( 62 | '--num-heads', type=int, default=8, help='number of attention heads in hidden layers' 63 | ) 64 | parser.add_argument( 65 | '--num-out-heads', type=int, default=1, help='number of attention heads in output layer' 66 | ) 67 | parser.add_argument('--dropout', type=float, default=0.6, help='dropout probability') 68 | parser.add_argument('--epochs', type=int, default=200, help='number of training epochs') 69 | parser.add_argument('--lr', type=float, default=0.005, help='learning rate') 70 | parser.add_argument('--weight-decay', type=float, default=5e-4, help='weight decay') 71 | args = parser.parse_args() 72 | print(args) 73 | train(args) 74 | 75 | 76 | if __name__ == '__main__': 77 | main() 78 | -------------------------------------------------------------------------------- /gnn/gcn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/gcn/__init__.py -------------------------------------------------------------------------------- /gnn/gcn/model.py: -------------------------------------------------------------------------------- 1 | """Graph Convolutional Network (GCN) 2 | 3 | 论文链接：https://arxiv.org/abs/1609.02907 4 | """ 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from dgl.nn import GraphConv 8 | 9 | 10 | class GCN(nn.Module): 11 | """两层GCN模型""" 12 | 13 | def __init__(self, in_dim, hidden_dim, out_dim): 14 | super().__init__() 15 | self.conv1 = GraphConv(in_dim, hidden_dim, activation=F.relu) 16 | self.conv2 = GraphConv(hidden_dim, out_dim) 17 | 18 | def forward(self, g, x): 19 | """ 20 | :param g: DGLGraph 21 | :param x: tensor(N, d_in) 输入顶点特征，N为g的顶点数 22 | :return: tensor(N, d_out) 输出顶点特征 23 | """ 24 | h = self.conv1(g, x) # (N, d_hid) 25 | h = self.conv2(g, h) # (N, d_out) 26 | return h 27 | -------------------------------------------------------------------------------- /gnn/gcn/readme.txt: -------------------------------------------------------------------------------- 1 | 顶点分类 2 | Cora 3 | python -m gnn.gcn.train --dataset=cora 4 | Test Accuracy 0.8070 5 | 6 | Citeseer 7 | python -m gnn.gcn.train --dataset=citeseer 8 | Test Accuracy 0.6950 9 | 10 | Pubmed 11 | python -m gnn.gcn.train --dataset=pubmed 12 | Test Accuracy 0.8010 13 | -------------------------------------------------------------------------------- /gnn/gcn/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import dgl 4 | import torch.nn as nn 5 | import torch.optim as optim 6 | 7 | from gnn.gcn.model import GCN 8 | from gnn.utils import load_citation_dataset, accuracy 9 | 10 | 11 | def train(args): 12 | data = load_citation_dataset(args.dataset) 13 | g = data[0] 14 | features = g.ndata['feat'] 15 | labels = g.ndata['label'] 16 | train_mask = g.ndata['train_mask'] 17 | val_mask = g.ndata['val_mask'] 18 | test_mask = g.ndata['test_mask'] 19 | 20 | # add self loop 21 | g = dgl.remove_self_loop(g) 22 | g = dgl.add_self_loop(g) 23 | 24 | model = GCN(features.shape[1], args.num_hidden, data.num_classes) 25 | criterion = nn.CrossEntropyLoss() 26 | optimizer = optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4) 27 | 28 | for epoch in range(args.epochs): 29 | model.train() 30 | logits = model(g, features) 31 | loss = criterion(logits[train_mask], labels[train_mask]) 32 | optimizer.zero_grad() 33 | loss.backward() 34 | optimizer.step() 35 | 36 | acc = accuracy(logits[val_mask], labels[val_mask]) 37 | print('Epoch {:05d} | Loss {:.4f} | ValAcc {:.4f}'.format(epoch, loss.item(), acc)) 38 | acc = accuracy(model(g, features)[test_mask], labels[test_mask]) 39 | print('Test Accuracy {:.4f}'.format(acc)) 40 | 41 | 42 | def main(): 43 | parser = argparse.ArgumentParser(description='GCN') 44 | parser.add_argument('--dataset', choices=['cora', 'citeseer', 'pubmed'], default='cora', help='dataset') 45 | parser.add_argument('--epochs', type=int, default=200, help='number of training epochs') 46 | parser.add_argument('--num-hidden', type=int, default=16, help='number of hidden units') 47 | args = parser.parse_args() 48 | train(args) 49 | 50 | 51 | if __name__ == '__main__': 52 | main() 53 | -------------------------------------------------------------------------------- /gnn/han/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/han/__init__.py -------------------------------------------------------------------------------- /gnn/han/model.py: -------------------------------------------------------------------------------- 1 | """Heterogeneous Graph Attention Network (HAN) 2 | 3 | 论文链接：https://arxiv.org/abs/1903.07293 4 | """ 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | 9 | from dgl.nn import GATConv 10 | 11 | 12 | class SemanticAttention(nn.Module): 13 | 14 | def __init__(self, in_dim, hidden_dim=128): 15 | """语义层次的注意力，将顶点基于不同元路径的嵌入组合为最终嵌入 16 | 17 | :param in_dim: 输入特征维数d_in，对应顶点层次注意力的输出维数 18 | :param hidden_dim: 语义层次隐含特征维数d_hid 19 | """ 20 | super().__init__() 21 | self.project = nn.Sequential( 22 | nn.Linear(in_dim, hidden_dim), # 论文公式(7)中的W和b 23 | nn.Tanh(), 24 | nn.Linear(hidden_dim, 1, bias=False) # 论文公式(7)中的q 25 | ) 26 | 27 | def forward(self, z): 28 | """ 29 | :param z: tensor(N, M, d_in) 顶点基于不同元路径的嵌入，N为顶点数，M为元路径个数 30 | :return: tensor(N, d_in) 顶点的最终嵌入 31 | """ 32 | w = self.project(z).mean(dim=0) # (N, M, d_hid) -> (N, M, 1) -> (M, 1) 33 | beta = torch.softmax(w, dim=0) # (M, 1) 34 | beta = beta.expand((z.shape[0],) + beta.shape) # (N, M, 1) 35 | z = (beta * z).sum(dim=1) # (N, d_in) 36 | return z 37 | 38 | 39 | class HANLayer(nn.Module): 40 | 41 | def __init__(self, num_metapaths, in_dim, out_dim, num_heads, dropout): 42 | """HAN层 43 | 44 | :param num_metapaths: int 元路径个数 45 | :param in_dim: int 输入特征维数 46 | :param out_dim: int 输出特征维数 47 | :param num_heads: int 注意力头数K 48 | :param dropout: float Dropout概率 49 | """ 50 | super().__init__() 51 | # 顶点层次的注意力，每个GAT层对应一个元路径 52 | self.gats = nn.ModuleList([ 53 | GATConv(in_dim, out_dim, num_heads, dropout, dropout, activation=F.elu) 54 | for _ in range(num_metapaths) 55 | ]) 56 | # 语义层次的注意力 57 | self.semantic_attention = SemanticAttention(in_dim=num_heads * out_dim) 58 | 59 | def forward(self, gs, h): 60 | """ 61 | :param gs: List[DGLGraph] 基于元路径的邻居组成的同构图 62 | :param h: tensor(N, d_in) 输入顶点特征 63 | :return: tensor(N, K*d_out) 输出顶点特征 64 | """ 65 | zp = [gat(g, h).flatten(start_dim=1) for gat, g in zip(self.gats, gs)] # 基于元路径的嵌入 66 | zp = torch.stack(zp, dim=1) # (N, M, K*d_out) 67 | z = self.semantic_attention(zp) # (N, K*d_out) 68 | return z 69 | 70 | 71 | class HAN(nn.Module): 72 | 73 | def __init__(self, num_metapaths, in_dim, hidden_dim, out_dim, num_heads, dropout): 74 | """HAN模型 75 | 76 | :param num_metapaths: int 元路径个数 77 | :param in_dim: int 输入特征维数 78 | :param hidden_dim: int 隐含特征维数 79 | :param out_dim: int 输出特征维数 80 | :param num_heads: int 注意力头数K 81 | :param dropout: float Dropout概率 82 | """ 83 | super().__init__() 84 | self.han = HANLayer(num_metapaths, in_dim, hidden_dim, num_heads, dropout) 85 | self.predict = nn.Linear(num_heads * hidden_dim, out_dim) 86 | 87 | def forward(self, gs, h): 88 | """ 89 | :param gs: List[DGLGraph] 基于元路径的邻居组成的同构图 90 | :param h: tensor(N, d_in) 输入顶点特征 91 | :return: tensor(N, d_out) 输出顶点嵌入 92 | """ 93 | h = self.han(gs, h) # (N, K*d_hid) 94 | out = self.predict(h) # (N, d_out) 95 | return out 96 | -------------------------------------------------------------------------------- /gnn/han/readme.txt: -------------------------------------------------------------------------------- 1 | 顶点分类 2 | ACM 3 | python -m gnn.han.train --dataset=acm --hetero --task=clf 4 | Test Micro-F1 0.9091 | Test Macro-F1 0.9114 5 | 6 | DBLP 7 | python -m gnn.han.train --dataset=dblp --hetero --task=clf 8 | Test Micro-F1 0.9247 | Test Macro-F1 0.9186 9 | 10 | IMDb 11 | python -m gnn.han.train --dataset=imdb --hetero --task=clf 12 | Test Micro-F1 0.5650 | Test Macro-F1 0.5623 13 | 14 | 顶点聚类 15 | ACM 16 | python -m gnn.han.train --dataset=acm --hetero --task=cluster 17 | Test NMI 0.7090 | Test ARI 0.7541 18 | 19 | DBLP 20 | python -m gnn.han.train --dataset=dblp --hetero --task=cluster 21 | Test NMI 0.7688 | Test ARI 0.8271 22 | 23 | IMDb 24 | python -m gnn.han.train --dataset=imdb --hetero --task=cluster 25 | Test NMI 0.1135 | Test ARI 0.1271 26 | -------------------------------------------------------------------------------- /gnn/han/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import dgl 4 | import torch 5 | import torch.nn.functional as F 6 | import torch.optim as optim 7 | 8 | from gnn.data import ACM3025Dataset, DBLP4057Dataset, IMDb5kDataset, ACMDataset, \ 9 | DBLPFourAreaDataset, IMDbDataset 10 | from gnn.han.model import HAN 11 | from gnn.utils import set_random_seed, micro_macro_f1_score, nmi_ari_score 12 | 13 | DATASET = { 14 | 'acm': ACM3025Dataset, 15 | 'dblp': DBLP4057Dataset, 16 | 'imdb': IMDb5kDataset 17 | } 18 | 19 | HETERO_DATASET = { 20 | 'acm': ACMDataset, 21 | 'dblp': DBLPFourAreaDataset, 22 | 'imdb': IMDbDataset 23 | } 24 | 25 | 26 | def train(args): 27 | set_random_seed(args.seed) 28 | if args.hetero: 29 | data = HETERO_DATASET[args.dataset]() 30 | g = data[0] 31 | gs = [dgl.metapath_reachable_graph(g, metapath) for metapath in data.metapaths] 32 | for i in range(len(gs)): 33 | gs[i] = dgl.add_self_loop(dgl.remove_self_loop(gs[i])) 34 | ntype = data.predict_ntype 35 | num_classes = data.num_classes 36 | features = g.nodes[ntype].data['feat'] 37 | labels = g.nodes[ntype].data['label'] 38 | train_mask = g.nodes[ntype].data['train_mask'] 39 | val_mask = g.nodes[ntype].data['val_mask'] 40 | test_mask = g.nodes[ntype].data['test_mask'] 41 | else: 42 | data = DATASET[args.dataset]() 43 | gs = data[0] 44 | num_classes = data.num_classes 45 | features = gs[0].ndata['feat'] 46 | labels = gs[0].ndata['label'] 47 | train_mask = gs[0].ndata['train_mask'] 48 | val_mask = gs[0].ndata['val_mask'] 49 | test_mask = gs[0].ndata['test_mask'] 50 | 51 | model = HAN( 52 | len(gs), features.shape[1], args.num_hidden, num_classes, args.num_heads, args.dropout 53 | ) 54 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 55 | 56 | score = micro_macro_f1_score if args.task == 'clf' else nmi_ari_score 57 | if args.task == 'clf': 58 | metrics = 'Epoch {:d} | Train Loss {:.4f} | Train Micro-F1 {:.4f} | Train Macro-F1 {:.4f}' \ 59 | ' | Val Micro-F1 {:.4f} | Val Macro-F1 {:.4f}' 60 | else: 61 | metrics = 'Epoch {:d} | Train Loss {:.4f} | Train NMI {:.4f} | Train ARI {:.4f}' \ 62 | ' | Val NMI {:.4f} | Val ARI {:.4f}' 63 | for epoch in range(args.epochs): 64 | model.train() 65 | logits = model(gs, features) 66 | loss = F.cross_entropy(logits[train_mask], labels[train_mask]) 67 | optimizer.zero_grad() 68 | loss.backward() 69 | optimizer.step() 70 | 71 | train_metrics = score(logits[train_mask], labels[train_mask]) 72 | val_metrics = score(logits[val_mask], labels[val_mask]) 73 | print(metrics.format(epoch, loss.item(), *train_metrics, *val_metrics)) 74 | 75 | test_metrics = evaluate(model, gs, features, labels, test_mask, score) 76 | if args.task == 'clf': 77 | print('Test Micro-F1 {:.4f} | Test Macro-F1 {:.4f}'.format(*test_metrics)) 78 | else: 79 | print('Test NMI {:.4f} | Test ARI {:.4f}'.format(*test_metrics)) 80 | 81 | 82 | def evaluate(model, gs, features, labels, mask, score): 83 | model.eval() 84 | with torch.no_grad(): 85 | logits = model(gs, features) 86 | return score(logits[mask], labels[mask]) 87 | 88 | 89 | def main(): 90 | parser = argparse.ArgumentParser('HAN Node Classification or Clustering') 91 | parser.add_argument('--seed', type=int, default=1, help='random seed') 92 | parser.add_argument('--dataset', choices=['acm', 'dblp', 'imdb'], default='acm', help='dataset') 93 | parser.add_argument( 94 | '--hetero', action='store_true', help='Use heterogeneous graph dataset' 95 | ) 96 | parser.add_argument('--num-hidden', type=int, default=8, help='number of hidden units') 97 | parser.add_argument( 98 | '--num-heads', type=int, default=8, help='number of attention heads in node-level attention' 99 | ) 100 | parser.add_argument('--dropout', type=float, default=0.6, help='dropout probability') 101 | parser.add_argument('--task', choices=['clf', 'cluster'], default='clf', help='training task') 102 | parser.add_argument('--epochs', type=int, default=200, help='number of training epochs') 103 | parser.add_argument('--lr', type=float, default=0.005, help='learning rate') 104 | parser.add_argument('--weight-decay', type=float, default=0.001, help='weight decay') 105 | args = parser.parse_args() 106 | print(args) 107 | train(args) 108 | 109 | 110 | if __name__ == '__main__': 111 | main() 112 | -------------------------------------------------------------------------------- /gnn/heco/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/heco/__init__.py -------------------------------------------------------------------------------- /gnn/heco/readme.md: -------------------------------------------------------------------------------- 1 | ## ACM 2 | `python -m gnn.heco.train --dataset=acm` 3 | 4 | ``` 5 | Micro-F1 0.8830 | Macro-F1 0.8850 | AUC 0.9700 6 | NMI 0.5954 | ARI 0.6138 7 | ``` 8 | 9 | 作者代码结果： 10 | ``` 11 | Loading 199th epoch 12 | [Classification] Macro-F1_mean: 0.8686 var: 0.0068 Micro-F1_mean: 0.8645 var: 0.0064 auc 0.9548 13 | [Classification] Macro-F1_mean: 0.8649 var: 0.0038 Micro-F1_mean: 0.8652 var: 0.0046 auc 0.9598 14 | [Classification] Macro-F1_mean: 0.8796 var: 0.0063 Micro-F1_mean: 0.8769 var: 0.0059 auc 0.9543 15 | NMI 0.5634 | ARI 0.5347 16 | ``` 17 | 18 | ## DBLP 19 | `python -m gnn.heco.train --dataset=dblp --feat-drop=0.4 --attn-drop=0.35 --tau=0.9` 20 | 21 | ``` 22 | Micro-F1 0.9070 | Macro-F1 0.9032 | AUC 0.9827 23 | NMI 0.6784 | ARI 0.7030 24 | ``` 25 | 26 | 作者代码结果： 27 | ``` 28 | Loading 196th epoch 29 | [Classification] Macro-F1_mean: 0.9122 var: 0.0032 Micro-F1_mean: 0.9188 var: 0.0033 auc 0.9816 30 | [Classification] Macro-F1_mean: 0.8970 var: 0.0029 Micro-F1_mean: 0.9001 var: 0.0039 auc 0.9796 31 | [Classification] Macro-F1_mean: 0.9092 var: 0.0027 Micro-F1_mean: 0.9174 var: 0.0026 auc 0.9844 32 | NMI 0.7086 | ARI 0.7645 33 | ``` 34 | 35 | ## Freebase 36 | `python -m gnn.heco.train --dataset=freebase --feat-drop=0.1 --attn-drop=0.3 --tau=0.5 --lr=0.001` 37 | 38 | ``` 39 | Micro-F1 0.5600 | Macro-F1 0.5404 | AUC 0.7307 40 | NMI 0.1773 | ARI 0.2003 41 | ``` 42 | 43 | 作者代码结果： 44 | ``` 45 | Loading 197th epoch 46 | [Classification] Macro-F1_mean: 0.5493 var: 0.0149 Micro-F1_mean: 0.5739 var: 0.0224 auc 0.7465 47 | [Classification] Macro-F1_mean: 0.6004 var: 0.0061 Micro-F1_mean: 0.6277 var: 0.0102 auc 0.7709 48 | [Classification] Macro-F1_mean: 0.5671 var: 0.0223 Micro-F1_mean: 0.5982 var: 0.0290 auc 0.7418 49 | NMI 0.1850 | ARI 0.2121 50 | ``` 51 | 52 | ## AMiner 53 | `python -m gnn.heco.train --dataset=aminer --feat-drop=0.5 --attn-drop=0.5 --tau=0.5 --lr=0.003` 54 | 55 | ``` 56 | Micro-F1 0.7350 | Macro-F1 0.6519 | AUC 0.8800 57 | NMI 0.3110 | ARI 0.2773 58 | ``` 59 | 60 | 作者代码结果： 61 | ``` 62 | Loading 199th epoch 63 | [Classification] Macro-F1_mean: 0.6728 var: 0.0067 Micro-F1_mean: 0.7421 var: 0.0071 auc 0.8743 64 | [Classification] Macro-F1_mean: 0.6851 var: 0.0032 Micro-F1_mean: 0.7569 var: 0.0034 auc 0.8874 65 | [Classification] Macro-F1_mean: 0.7232 var: 0.0026 Micro-F1_mean: 0.7927 var: 0.0023 auc 0.9038 66 | NMI 0.3170 | ARI 0.2810 67 | ``` 68 | 69 | 70 | ## 踩坑记录 71 | 一开始自己实现的模型性能比作者提供的代码差很多 72 | 73 | 花了两天时间逐个部分对比差异，结果发现该模型非常不稳定，对数据、超参数等各种细节非常敏感： 74 | 75 | * 基于元路径的邻居图的邻接矩阵是否归一化 76 | * 将作者自己实现的GCN改为dgl的GraphConv、归一化方式由both改为right，性能提升 77 | * 将作者自己实现的GAT改为dgl的GATConv，性能下降（主要区别在于attn_drop的使用方式） 78 | * 参数初始化（使用nn.init.xavier_normal_会提升） 79 | * 正样本采样（np.argsort可以 torch.argsort或torch.topk就不行） 80 | 81 | 都对结果有很大影响（尤其是顶点聚类）！ 82 | 83 | 经过修改，基本达到与作者代码相近的结果 84 | -------------------------------------------------------------------------------- /gnn/hetgnn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/hetgnn/__init__.py -------------------------------------------------------------------------------- /gnn/hetgnn/eval.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import pickle 3 | 4 | import torch 5 | from sklearn.cluster import KMeans 6 | from sklearn.linear_model import LogisticRegression 7 | from sklearn.metrics import f1_score, normalized_mutual_info_score, adjusted_rand_score 8 | from sklearn.model_selection import train_test_split 9 | 10 | from gnn.data import AMiner2Dataset 11 | 12 | 13 | def load_data(node_embed_path): 14 | data = AMiner2Dataset() 15 | g = data[0] 16 | idx = torch.nonzero(g.nodes['author'].data['label'] != -1, as_tuple=True)[0] 17 | labels = g.nodes['author'].data['label'][idx].numpy() 18 | 19 | with open(node_embed_path, 'rb') as f: 20 | node_embeds = pickle.load(f) 21 | return node_embeds['author'][idx].numpy(), labels, data.num_classes 22 | 23 | 24 | def node_clf(args): 25 | embeds, labels, _ = load_data(args.node_embed_path) 26 | for train_size in (0.1, 0.3): 27 | X_train, X_test, y_train, y_test = train_test_split( 28 | embeds, labels, train_size=train_size, random_state=args.seed 29 | ) 30 | clf = LogisticRegression(random_state=args.seed, max_iter=500) 31 | clf.fit(X_train, y_train) 32 | y_pred = clf.predict(X_test) 33 | micro_f1 = f1_score(y_test, y_pred, average='micro') 34 | macro_f1 = f1_score(y_test, y_pred, average='macro') 35 | print('Train size {:.0%} | Test Micro-F1 {:.4f} | Test Macro-F1 {:.4f}'.format( 36 | train_size, micro_f1, macro_f1 37 | )) 38 | 39 | 40 | def node_cluster(args): 41 | embeds, labels, num_classes = load_data(args.node_embed_path) 42 | pred = KMeans(num_classes, random_state=args.seed).fit_predict(embeds) 43 | nmi = normalized_mutual_info_score(labels, pred) 44 | ari = adjusted_rand_score(labels, pred) 45 | print('NMI {:.4f} | ARI {:.4f}'.format(nmi, ari)) 46 | 47 | 48 | def main(): 49 | parser = argparse.ArgumentParser(description='HetGNN node classification or clustering') 50 | parser.add_argument('--seed', type=int, default=10, help='random seed') 51 | parser.add_argument('--task', choices=['clf', 'cluster'], default='clf', help='training task') 52 | parser.add_argument('node_embed_path', help='path to saved node embeddings') 53 | args = parser.parse_args() 54 | if args.task == 'clf': 55 | node_clf(args) 56 | else: 57 | node_cluster(args) 58 | 59 | 60 | if __name__ == '__main__': 61 | main() 62 | -------------------------------------------------------------------------------- /gnn/hetgnn/model.py: -------------------------------------------------------------------------------- 1 | """Heterogeneous Graph Neural Network (HetGNN) 2 | 3 | 论文链接：https://dl.acm.org/doi/pdf/10.1145/3292500.3330961 4 | """ 5 | import dgl.function as fn 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | 11 | class ContentAggregation(nn.Module): 12 | 13 | def __init__(self, in_dim, hidden_dim): 14 | """异构内容嵌入模块，针对一种顶点类型，将该类型顶点的多个输入特征编码为一个向量 15 | 16 | :param in_dim: int 输入特征维数 17 | :param hidden_dim: int 内容嵌入维数 18 | """ 19 | super().__init__() 20 | self.lstm = nn.LSTM(in_dim, hidden_dim // 2, batch_first=True, bidirectional=True) 21 | 22 | def forward(self, feats): 23 | """ 24 | :param feats: tensor(N, C, d_in) 输入特征列表，N为batch大小，C为输入特征个数 25 | :return: tensor(N, d_hid) 顶点的异构内容嵌入向量 26 | """ 27 | out, _ = self.lstm(feats) # (N, C, d_hid) 28 | return torch.mean(out, dim=1) 29 | 30 | 31 | class NeighborAggregation(nn.Module): 32 | 33 | def __init__(self, emb_dim): 34 | """邻居聚集模块，针对一种邻居类型t，将一个顶点的所有t类型邻居的内容嵌入向量聚集为一个向量 35 | 36 | :param emb_dim: int 内容嵌入维数 37 | """ 38 | super().__init__() 39 | self.lstm = nn.LSTM(emb_dim, emb_dim // 2, batch_first=True, bidirectional=True) 40 | 41 | def forward(self, embeds): 42 | """ 43 | :param embeds: tensor(N, Nt, d) 邻居的内容嵌入，N为batch大小，Nt为每个顶点的邻居个数 44 | :return: tensor(N, d) 顶点的t类型邻居聚集嵌入 45 | """ 46 | out, _ = self.lstm(embeds) # (N, Nt, d) 47 | return torch.mean(out, dim=1) 48 | 49 | 50 | class TypesCombination(nn.Module): 51 | 52 | def __init__(self, emb_dim): 53 | """类型组合模块，针对一种顶点类型，将该类型顶点的所有类型的邻居聚集嵌入组合为一个向量 54 | 55 | :param emb_dim: int 邻居嵌入维数 56 | """ 57 | super().__init__() 58 | self.attn = nn.Parameter(torch.ones(1, 2 * emb_dim)) 59 | self.leaky_relu = nn.LeakyReLU() 60 | 61 | def forward(self, content_embed, neighbor_embeds): 62 | """ 63 | :param content_embed: tensor(N, d) 内容嵌入，N为batch大小 64 | :param neighbor_embeds: tensor(A, N, d) 邻居嵌入，A为邻居类型数 65 | :return: tensor(N, d) 最终嵌入 66 | """ 67 | neighbor_embeds = torch.cat( 68 | [content_embed.unsqueeze(0), neighbor_embeds], dim=0 69 | ) # (A+1, N, d) 70 | cat_embeds = torch.cat( 71 | [content_embed.repeat(neighbor_embeds.shape[0], 1, 1), neighbor_embeds], dim=-1 72 | ) # (A+1, N, 2d) 73 | attn_scores = self.leaky_relu((self.attn * cat_embeds).sum(dim=-1, keepdim=True)) 74 | attn_scores = F.softmax(attn_scores, dim=0) # (A+1, N, 1) 75 | out = (attn_scores * neighbor_embeds).sum(dim=0) # (N, d) 76 | return out 77 | 78 | 79 | class HetGNN(nn.Module): 80 | 81 | def __init__(self, in_dim, hidden_dim, ntypes): 82 | """HetGNN模型 83 | 84 | :param in_dim: int 输入特征维数 85 | :param hidden_dim: int 隐含特征维数 86 | :param ntypes: List[str] 顶点类型列表 87 | """ 88 | super().__init__() 89 | self.content_aggs = nn.ModuleDict({ 90 | ntype: ContentAggregation(in_dim, hidden_dim) for ntype in ntypes 91 | }) 92 | self.neighbor_aggs = nn.ModuleDict({ 93 | ntype: NeighborAggregation(hidden_dim) for ntype in ntypes 94 | }) 95 | self.combs = nn.ModuleDict({ 96 | ntype: TypesCombination(hidden_dim) for ntype in ntypes 97 | }) 98 | 99 | def forward(self, g, feats): 100 | """ 101 | :param g: DGLGraph 异构图 102 | :param feats: Dict[str, tensor(N_i, C_i, d_in)] 顶点类型到输入特征的映射 103 | :return: Dict[str, tensor(N_i, d_hid)] 顶点类型到输出特征的映射 104 | """ 105 | with g.local_scope(): 106 | # 1.异构内容聚集 107 | for ntype in g.ntypes: 108 | g.nodes[ntype].data['c'] = self.content_aggs[ntype](feats[ntype]) # (N_i, d_hid) 109 | 110 | # 2.同类型邻居聚集 111 | neighbor_embeds = {} 112 | for dt in g.ntypes: 113 | tmp = [] 114 | for st in g.ntypes: 115 | # dt类型所有顶点的st类型邻居个数必须全部相同 116 | g.multi_update_all({f'{st}-{dt}': (fn.copy_u('c', 'm'), stack_reducer)}, 'sum') 117 | tmp.append(self.neighbor_aggs[st](g.nodes[dt].data.pop('nc'))) 118 | neighbor_embeds[dt] = torch.stack(tmp) # (A, N_dt, d_hid) 119 | 120 | # 3.类型组合 121 | out = { 122 | ntype: self.combs[ntype](g.nodes[ntype].data['c'], neighbor_embeds[ntype]) 123 | for ntype in g.ntypes 124 | } 125 | return out 126 | 127 | def calc_score(self, g, h): 128 | """计算图中每一条边的得分 s(u, v)=h(u)^T h(v) 129 | 130 | :param g: DGLGraph 异构图 131 | :param h: Dict[str, tensor(N_i, d)] 顶点类型到顶点嵌入的映射 132 | :return: tensor(A*E) 所有边的得分 133 | """ 134 | with g.local_scope(): 135 | g.ndata['h'] = h 136 | for etype in g.etypes: 137 | g.apply_edges(fn.u_dot_v('h', 'h', 's'), etype=etype) 138 | return torch.cat(list(g.edata['s'].values())).squeeze(dim=-1) # (A*E,) 139 | 140 | 141 | def stack_reducer(nodes): 142 | return {'nc': nodes.mailbox['m']} 143 | -------------------------------------------------------------------------------- /gnn/hetgnn/preprocess.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | from dgl.data.utils import save_graphs, save_info 5 | 6 | from gnn.hetgnn.utils import load_data 7 | 8 | 9 | def preprocess(args): 10 | g, feats = load_data(args.neighbor_path, args.pretrained_node_embed_path) 11 | save_graphs(os.path.join(args.save_path, 'neighbor_graph.bin'), [g]) 12 | save_info(os.path.join(args.save_path, 'in_feats.pkl'), feats) 13 | print('Neighbor graph and input features saved to', args.save_path) 14 | 15 | 16 | def main(): 17 | parser = argparse.ArgumentParser(description='HetGNN preprocessing') 18 | parser.add_argument('neighbor_path', help='path to neighbor file generated by random walk') 19 | parser.add_argument('pretrained_node_embed_path', help='path to pretrained node embeddings') 20 | parser.add_argument( 21 | 'save_path', help='path to save preprocessed neighbor graph and input features' 22 | ) 23 | args = parser.parse_args() 24 | preprocess(args) 25 | 26 | 27 | if __name__ == '__main__': 28 | main() 29 | -------------------------------------------------------------------------------- /gnn/hetgnn/random_walk.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from gnn.data import AMiner2Dataset 4 | from gnn.utils import metapath_random_walk 5 | 6 | 7 | def main(): 8 | parser = argparse.ArgumentParser(description='Random walk with restart') 9 | parser.add_argument('--num-walks', type=int, default=10, help='number of walks for each node') 10 | parser.add_argument('--walk-length', type=int, default=30, help='times to repeat metapath') 11 | parser.add_argument('output_file', help='output filename') 12 | args = parser.parse_args() 13 | 14 | data = AMiner2Dataset() 15 | g = data[0] 16 | 17 | # author_0特殊处理，添加假边 18 | g.add_edges([0] * 10, list(range(10)), etype='write') 19 | g.add_edges(list(range(10)), [0] * 10, etype='write_rev') 20 | 21 | metapaths = { 22 | 'author': ['write', 'publish', 'publish_rev', 'write_rev'], # APVPA 23 | 'paper': ['write_rev', 'write', 'publish', 'publish_rev'], # PAPVP 24 | 'venue': ['publish_rev', 'write_rev', 'write', 'publish'] # VPAPV 25 | } 26 | metapath_random_walk(g, metapaths, args.num_walks, args.walk_length, args.output_file) 27 | 28 | 29 | if __name__ == '__main__': 30 | main() 31 | -------------------------------------------------------------------------------- /gnn/hetgnn/readme.md: -------------------------------------------------------------------------------- 1 | # 运行代码 2 | ## 1.随机游走 3 | `python -m gnn.hetgnn.random_walk D:\dgldata\aminer2\corpus.txt` 4 | 5 | ## 2.预训练顶点嵌入 6 | `python -m gnn.metapath2vec.train_word2vec --size=128 D:\dgldata\aminer2\corpus.txt D:\dgldata\aminer2\aminer2.model` 7 | 8 | ## 3.预处理，生成邻居图和输入特征 9 | `python -m gnn.hetgnn.preprocess D:\dgldata\aminer2\corpus.txt D:\dgldata\aminer2\aminer2.model D:\dgldata\aminer2\` 10 | 11 | ## 4.训练模型（无监督），学习顶点嵌入 12 | `python -m gnn.hetgnn.train D:\dgldata\aminer2\ D:\dgldata\aminer2\node_embed.pkl` 13 | 14 | ## 5.评估顶点嵌入 15 | ### 顶点分类 16 | `python -m gnn.hetgnn.eval --task=clf D:\dgldata\aminer2\node_embed.pkl` 17 | ``` 18 | Train size 10% | Test Micro-F1 0.9588 | Test Macro-F1 0.9573 19 | Train size 30% | Test Micro-F1 0.9657 | Test Macro-F1 0.9647 20 | ``` 21 | 22 | ### 顶点聚类 23 | `python -m gnn.hetgnn.eval --task=cluster D:\dgldata\aminer2\node_embed.pkl` 24 | ``` 25 | NMI 0.8099 | ARI 0.8611 26 | ``` 27 | 28 | # 数据集 29 | 作者提供的数据集（只有论文中的Academic II数据集）： 30 | * https://github.com/chuxuzhang/KDD2019_HetGNN/tree/master/data/academic 31 | * https://drive.google.com/file/d/1N6GWsniacaT-L0GPXpi1D3gM2LVOih-A/view?usp=sharing 32 | 33 | 二者基本相同，但GitHub上的数据包含随机游走结果、生成的训练样本等，都没有原始数据 34 | 35 | ## 图结构 36 | * 顶点：28646 author, 21044 paper, 18 venue （学者顶点没有0） 37 | * 边：A-P, P-P, P-V 38 | 39 | ## 文件说明 40 | ### A-P边 41 | * a_p_list_train.txt, a_p_list_test.txt 格式：aid:pid,pid,... 42 | * p_a_list_train.txt, p_a_list_test.txt 格式：pid:aid,aid,... （反向边，与上面两个文件完全相同） 43 | 44 | ### P-P边 45 | * p_p_cite_list_train.txt, p_p_cite_list_test.txt 格式：pid:pid,pid,... 46 | * p_p_citation_list.txt不是两者的并集，无用 47 | 48 | ### P-V边 49 | * p_v.txt 格式：pid,vid （完整，0<=pid<=21043, 0<=vid<=17） 50 | * v_p_list_train.txt 格式：vid:pid,pid,... （不完整） 51 | 52 | ### author标签 53 | a_class_train.txt, a_class_test.txt 格式：aid,class （不完整，11360个aid, 0<=class<=3） 54 | 55 | ### paper特征 56 | * p_title_embed.txt 标题词向量，格式：（第二行开始）pid x1 ... x128 （完整，0<=pid<=21043） 57 | * p_abstract_embed.txt 摘要词向量，格式同上 58 | * p_time.txt 论文年份，格式：pid\ttime （完整，time是论文年份-2005，见作者代码raw_data_process.py） 59 | 60 | ### 随机游走结果 61 | * het_random_walk.txt 格式：{a|p|v}_{id} × n （不完整，出现的顶点集合同node_net_embedding.txt） 62 | * het_neigh_train.txt 格式：{a|p|v}\_{id}:{a|p|v}\_{id},... （不完整，出现的顶点集合同node_net_embedding.txt） 63 | 64 | ### 顶点嵌入 65 | node_net_embedding.txt 随机游走结果+word2vec 格式：{a|p|v}_{id} x1 ... x128 66 | （不完整，数量：22821 author, 15930 paper, 18 venue） 67 | 68 | ### 生成的训练样本 69 | * A-A: a_a_list_train.txt, a_a_list_test.txt 格式：aid, aid, {0|1} （0/1的含义未知） 70 | * A-P: a_p_cite_list_train.txt, a_p_cite_list_test.txt 格式：aid, pid, {0|1} （0/1的含义未知） 71 | * A-V: a_v_list_train.txt, a_v_list_test.txt 格式：aid:vid,vid,..., 72 | 73 | # 作者代码 74 | 作者提供的代码：https://github.com/chuxuzhang/KDD2019_HetGNN/tree/master/code 75 | 76 | （~~又臭又长~~） 77 | * raw_data_process.py 将原始数据处理成提供的格式（没有原始数据，因此无用） 78 | * input_data_process.py 随机游走和生成三种训练样本边(A-A, A-P, A-V) 79 | * 依次从每一个学者顶点出发，游走固定长度L=30，每个顶点重复n=10次 80 | * 随机游走策略：如果当前顶点为A或V则随机选择一个P邻居；如果当前顶点为P则从A+P+V邻居中随机选择一个 81 | * DeepWalk.py 随机游走结果+word2vec → 预训练的顶点嵌入(node_net_embedding.txt) 82 | * data_generator.py 读取原始数据 83 | * paper顶点特征 84 | * p_title_embed 预训练的标题词向量 85 | * p_abstract_embed 预训练的摘要词向量 86 | * p_net_embed 预训练的顶点嵌入 87 | * p_v_net_embed 论文所属期刊的v_net_embed 88 | * p_a_net_embed 论文作者的a_net_embed的平均 89 | * p_ref_net_embed 论文引用的论文的p_net_embed的平均 90 | * author顶点特征 91 | * a_net_embed 预训练的顶点嵌入 92 | * a_text_embed 3篇论文的p_abstract_embed的拼接（不足3篇的用最后一篇补齐） 93 | * venue顶点特征 94 | * v_net_embed 预训练的顶点嵌入 95 | * v_text_embed 5篇论文的p_abstract_embed的拼接（不足5篇的用最后一篇补齐） 96 | * het_walk_restart() 随机游走生成邻居(het_neigh_train.txt) 97 | * 每种顶点都选择邻居中出现次数最多的10个author、10个paper、3个venue 98 | * input_data_process.py中的随机游走是为了预训练顶点嵌入生成语料库，与这里的随机游走不同 99 | * 负采样：sample_het_walk_triple() 100 | * args.py 解析命令行参数 101 | * tools.py 核心模型HetAgg（你管这叫tools？） 102 | * HetGNN.py 整体模型+训练 103 | -------------------------------------------------------------------------------- /gnn/hetgnn/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import pickle 4 | 5 | import torch 6 | import torch.nn.functional as F 7 | import torch.optim as optim 8 | from dgl.data.utils import load_graphs, load_info 9 | 10 | from gnn.hetgnn.model import HetGNN 11 | from gnn.hetgnn.utils import construct_neg_graph 12 | from gnn.utils import set_random_seed, RatioNegativeSampler 13 | 14 | 15 | def train(args): 16 | set_random_seed(args.seed) 17 | g = load_graphs(os.path.join(args.data_path, 'neighbor_graph.bin'))[0][0] 18 | feats = load_info(os.path.join(args.data_path, 'in_feats.pkl')) 19 | 20 | model = HetGNN(feats['author'].shape[-1], args.num_hidden, g.ntypes) 21 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 22 | neg_sampler = RatioNegativeSampler() 23 | for epoch in range(args.epochs): 24 | model.train() 25 | embeds = model(g, feats) 26 | score = model.calc_score(g, embeds) 27 | neg_g = construct_neg_graph(g, neg_sampler) 28 | neg_score = model.calc_score(neg_g, embeds) 29 | logits = torch.cat([score, neg_score]) # (2A*E,) 30 | labels = torch.cat([torch.ones(score.shape[0]), torch.zeros(neg_score.shape[0])]) 31 | loss = F.binary_cross_entropy_with_logits(logits, labels) 32 | 33 | optimizer.zero_grad() 34 | loss.backward() 35 | optimizer.step() 36 | print('Epoch {:d} | Loss {:.4f}'.format(epoch, loss.item())) 37 | with torch.no_grad(): 38 | final_embeds = model(g, feats) 39 | with open(args.save_node_embed_path, 'wb') as f: 40 | pickle.dump(final_embeds, f) 41 | print('Final node embeddings saved to', args.save_node_embed_path) 42 | 43 | 44 | def main(): 45 | parser = argparse.ArgumentParser(description='HetGNN unsupervised training') 46 | parser.add_argument('--seed', type=int, default=10, help='random seed') 47 | parser.add_argument('--num-hidden', type=int, default=128, help='number of hidden units') 48 | parser.add_argument('--epochs', type=int, default=100, help='number of training epochs') 49 | parser.add_argument('--lr', type=float, default=0.001, help='learning rate') 50 | parser.add_argument('--weight-decay', type=float, default=0.0, help='weight decay') 51 | parser.add_argument('data_path', help='path to preprocessed data') 52 | parser.add_argument('save_node_embed_path', help='path to save learned final node embeddings') 53 | args = parser.parse_args() 54 | print(args) 55 | train(args) 56 | 57 | 58 | if __name__ == '__main__': 59 | main() 60 | -------------------------------------------------------------------------------- /gnn/hetgnn/utils.py: -------------------------------------------------------------------------------- 1 | from collections import Counter 2 | from itertools import chain 3 | 4 | import dgl 5 | import dgl.function as fn 6 | import torch 7 | from gensim.models import Word2Vec 8 | 9 | from gnn.data import AMiner2Dataset 10 | 11 | 12 | def load_data(neighbor_path, node_embed_path): 13 | data = AMiner2Dataset() 14 | g = data[0] 15 | # author_0特殊处理，添加假边 16 | g.add_edges(0, 0, etype='write') 17 | g.add_edges(0, 0, etype='write_rev') 18 | 19 | print('Loading pretrained node embeddings...') 20 | load_pretrained_node_embed(g, node_embed_path) 21 | print('Propagating input features...') 22 | feats = propagate_feature(g) 23 | print('Constructing neighbor graph...') 24 | ng = construct_neighbor_graph(g, neighbor_path, {'author': 10, 'paper': 10, 'venue': 3}) 25 | return ng, feats 26 | 27 | 28 | def load_pretrained_node_embed(g, node_embed_path): 29 | model = Word2Vec.load(node_embed_path) 30 | for ntype in g.ntypes: 31 | g.nodes[ntype].data['net_embed'] = torch.from_numpy( 32 | model.wv[[f'{ntype}_{i}' for i in range(g.num_nodes(ntype))]] 33 | ) 34 | 35 | 36 | def propagate_feature(g): 37 | with g.local_scope(): 38 | g.multi_update_all({ 39 | 'write': (fn.copy_u('net_embed', 'm'), fn.mean('m', 'a_net_embed')), 40 | 'publish_rev': (fn.copy_u('net_embed', 'm'), fn.mean('m', 'v_net_embed')) 41 | }, 'sum') 42 | paper_feats = torch.stack([ 43 | g.nodes['paper'].data[k] for k in 44 | ('abstract_embed', 'title_embed', 'net_embed', 'v_net_embed', 'a_net_embed') 45 | ], dim=1) # (N_p, 5, d) 46 | 47 | ap = find_neighbors(g, 'write_rev', 3).view(1, -1) # (1, 3N_a) 48 | ap_abstract_embed = g.nodes['paper'].data['abstract_embed'][ap] \ 49 | .view(g.num_nodes('author'), 3, -1) # (N_a, 3, d) 50 | author_feats = torch.cat([ 51 | g.nodes['author'].data['net_embed'].unsqueeze(dim=1), ap_abstract_embed 52 | ], dim=1) # (N_a, 4, d) 53 | 54 | vp = find_neighbors(g, 'publish', 5).view(1, -1) # (1, 5N_v) 55 | vp_abstract_embed = g.nodes['paper'].data['abstract_embed'][vp] \ 56 | .view(g.num_nodes('venue'), 5, -1) # (N_v, 5, d) 57 | venue_feats = torch.cat([ 58 | g.nodes['venue'].data['net_embed'].unsqueeze(dim=1), vp_abstract_embed 59 | ], dim=1) # (N_v, 6, d) 60 | 61 | return {'author': author_feats, 'paper': paper_feats, 'venue': venue_feats} 62 | 63 | 64 | def find_neighbors(g, etype, n): 65 | num_nodes = g.num_nodes(g.to_canonical_etype(etype)[2]) 66 | u, v = g.in_edges(torch.arange(num_nodes), etype=etype) 67 | neighbors = [[] for _ in range(num_nodes)] 68 | for i in range(len(v)): 69 | neighbors[v[i].item()].append(u[i].item()) 70 | for v in range(num_nodes): 71 | if len(neighbors[v]) < n: 72 | neighbors[v] += [neighbors[v][-1]] * (n - len(neighbors[v])) 73 | elif len(neighbors[v]) > n: 74 | neighbors[v] = neighbors[v][:n] 75 | return torch.tensor(neighbors) # (N_dst, n) 76 | 77 | 78 | def construct_neighbor_graph(g, neighbor_path, neighbor_size): 79 | counts = { 80 | f'{stype}-{dtype}': [Counter() for _ in range(g.num_nodes(dtype))] 81 | for stype in g.ntypes for dtype in g.ntypes 82 | } 83 | with open(neighbor_path) as f: 84 | for line in f: 85 | center, neighbors = line.strip().split(' ', 1) 86 | dtype, v = parse_node_name(center) 87 | neighbors = neighbors.split(' ') 88 | for n in neighbors: 89 | stype, u = parse_node_name(n) 90 | counts[f'{stype}-{dtype}'][v][u] += 1 91 | 92 | edges = {} 93 | for dtype in g.ntypes: 94 | for stype in g.ntypes: 95 | etype = f'{stype}-{dtype}' 96 | edges[(stype, etype, dtype)] = (torch.tensor(list(chain.from_iterable( 97 | (u for u, _ in counts[etype][v].most_common(neighbor_size[stype])) 98 | for v in range(g.num_nodes(dtype)) 99 | ))), torch.arange(g.num_nodes(dtype)).repeat_interleave(neighbor_size[stype])) 100 | 101 | ng = dgl.heterograph(edges, {ntype: g.num_nodes(ntype) for ntype in g.ntypes}) 102 | for dtype in ng.ntypes: 103 | for stype in ng.ntypes: 104 | assert torch.all(ng.in_degrees(etype=f'{stype}-{dtype}') == neighbor_size[stype]), \ 105 | f'not all in-degrees of edge type {stype}-{dtype} are {neighbor_size[stype]}' 106 | return ng 107 | 108 | 109 | def parse_node_name(node): 110 | ntype, nid = node.split('_') 111 | return ntype, int(nid) 112 | 113 | 114 | def construct_neg_graph(g, neg_sampler): 115 | return dgl.heterograph( 116 | neg_sampler(g, {etype: torch.arange(g.num_edges(etype)) for etype in g.etypes}), 117 | {ntype: g.num_nodes(ntype) for ntype in g.ntypes} 118 | ) 119 | -------------------------------------------------------------------------------- /gnn/hgconv/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/hgconv/__init__.py -------------------------------------------------------------------------------- /gnn/hgconv/readme.txt: -------------------------------------------------------------------------------- 1 | 顶点分类 2 | ACM 3 | python -m gnn.hgconv.train --dataset=acm --task=clf --epochs=20 4 | Test Micro-F1 0.9180 | Test Macro-F1 0.9192 5 | 6 | IMDb 7 | python -m gnn.hgconv.train --dataset=imdb --task=clf 8 | Test Micro-F1 0.5506 | Test Macro-F1 0.5471 9 | 10 | 顶点聚类 11 | ACM 12 | python -m gnn.hgconv.train --dataset=acm --task=cluster 13 | Test NMI 0.7099 | Test ARI 0.7600 14 | 15 | IMDb 16 | python -m gnn.hgconv.train --dataset=imdb --task=cluster --epochs=10 17 | Test NMI 0.1150 | Test ARI 0.0907 18 | -------------------------------------------------------------------------------- /gnn/hgconv/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import warnings 3 | 4 | import torch 5 | import torch.nn.functional as F 6 | import torch.optim as optim 7 | 8 | from gnn.data import ACMDataset, IMDbDataset 9 | from gnn.hgconv.model import HGConv 10 | from gnn.utils import set_random_seed, micro_macro_f1_score, nmi_ari_score 11 | 12 | DATASET = { 13 | 'acm': ACMDataset, 14 | 'imdb': IMDbDataset 15 | } 16 | 17 | 18 | def train(args): 19 | set_random_seed(args.seed) 20 | data = DATASET[args.dataset]() 21 | g = data[0] 22 | predict_ntype = data.predict_ntype 23 | features = {ntype: g.nodes[ntype].data['feat'] for ntype in g.ntypes} 24 | labels = g.nodes[predict_ntype].data['label'] 25 | train_mask = g.nodes[predict_ntype].data['train_mask'] 26 | val_mask = g.nodes[predict_ntype].data['val_mask'] 27 | test_mask = g.nodes[predict_ntype].data['test_mask'] 28 | 29 | model = HGConv( 30 | {ntype: g.nodes[ntype].data['feat'].shape[1] for ntype in g.ntypes}, 31 | args.num_hidden, data.num_classes, args.num_heads, g.ntypes, g.canonical_etypes, 32 | predict_ntype, args.num_layers, args.dropout, args.residual 33 | ) 34 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 35 | 36 | score = micro_macro_f1_score if args.task == 'clf' else nmi_ari_score 37 | if args.task == 'clf': 38 | metrics = 'Epoch {:d} | Train Loss {:.4f} | Train Micro-F1 {:.4f} | Train Macro-F1 {:.4f}' \ 39 | ' | Val Micro-F1 {:.4f} | Val Macro-F1 {:.4f}' \ 40 | ' | Test Micro-F1 {:.4f} | Test Macro-F1 {:.4f}' 41 | else: 42 | metrics = 'Epoch {:d} | Train Loss {:.4f} | Train NMI {:.4f} | Train ARI {:.4f}' \ 43 | ' | Val NMI {:.4f} | Val ARI {:.4f}' \ 44 | ' | Test NMI {:.4f} | Test ARI {:.4f}' 45 | warnings.filterwarnings('ignore', 'Setting attributes on ParameterDict is not supported') 46 | for epoch in range(args.epochs): 47 | model.train() 48 | logits = model(g, features) 49 | loss = F.cross_entropy(logits[train_mask], labels[train_mask]) 50 | optimizer.zero_grad() 51 | loss.backward() 52 | optimizer.step() 53 | 54 | train_metrics = score(logits[train_mask], labels[train_mask]) 55 | val_metrics = evaluate(model, g, features, labels, val_mask, score) 56 | test_metrics = evaluate(model, g, features, labels, test_mask, score) 57 | print(metrics.format(epoch, loss.item(), *train_metrics, *val_metrics, *test_metrics)) 58 | 59 | test_metrics = evaluate(model, g, features, labels, test_mask, score) 60 | if args.task == 'clf': 61 | print('Test Micro-F1 {:.4f} | Test Macro-F1 {:.4f}'.format(*test_metrics)) 62 | else: 63 | print('Test NMI {:.4f} | Test ARI {:.4f}'.format(*test_metrics)) 64 | 65 | 66 | @torch.no_grad() 67 | def evaluate(model, g, features, labels, mask, score): 68 | model.eval() 69 | logits = model(g, features) 70 | return score(logits[mask], labels[mask]) 71 | 72 | 73 | def main(): 74 | parser = argparse.ArgumentParser(description='HGConv Node Classification or Clustering') 75 | parser.add_argument('--seed', type=int, default=1, help='random seed') 76 | parser.add_argument('--dataset', choices=['acm', 'imdb'], default='acm', help='dataset') 77 | parser.add_argument('--num-hidden', type=int, default=64, help='number of hidden units') 78 | parser.add_argument('--num-heads', type=int, default=8, help='number of attention heads') 79 | parser.add_argument('--num-layers', type=int, default=2, help='number of layers') 80 | parser.add_argument('--dropout', type=float, default=0.5, help='dropout probability') 81 | parser.add_argument( 82 | '--no-residual', action='store_false', help='no residual connection', dest='residual' 83 | ) 84 | parser.add_argument('--task', choices=['clf', 'cluster'], default='clf', help='training task') 85 | parser.add_argument('--epochs', type=int, default=50, help='number of training epochs') 86 | parser.add_argument('--lr', type=float, default=0.001, help='learning rate') 87 | parser.add_argument('--weight-decay', type=float, default=0.001, help='weight decay') 88 | args = parser.parse_args() 89 | print(args) 90 | train(args) 91 | 92 | 93 | if __name__ == '__main__': 94 | main() 95 | -------------------------------------------------------------------------------- /gnn/hgt/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/hgt/__init__.py -------------------------------------------------------------------------------- /gnn/hgt/readme.txt: -------------------------------------------------------------------------------- 1 | 顶点分类 2 | ACM 3 | python -m gnn.hgt.train --dataset=acm 4 | Test Micro-F1 0.9137 | Test Macro-F1 0.9145 5 | 6 | IMDb 7 | python -m gnn.hgt.train --dataset=imdb 8 | Test Micro-F1 0.5702 | Test Macro-F1 0.5690 9 | 10 | 11 | 未使用RTE，需要将k改为边特征 12 | -------------------------------------------------------------------------------- /gnn/hgt/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import warnings 3 | 4 | import torch 5 | import torch.nn.functional as F 6 | import torch.optim as optim 7 | 8 | from gnn.data import ACMDataset, IMDbDataset 9 | from gnn.hgt.model import HGT 10 | from gnn.utils import set_random_seed, micro_macro_f1_score 11 | 12 | DATASET = { 13 | 'acm': ACMDataset, 14 | 'imdb': IMDbDataset 15 | } 16 | 17 | 18 | def train(args): 19 | set_random_seed(args.seed) 20 | data = DATASET[args.dataset]() 21 | g = data[0] 22 | predict_ntype = data.predict_ntype 23 | features = {ntype: g.nodes[ntype].data['feat'] for ntype in g.ntypes} 24 | labels = g.nodes[predict_ntype].data['label'] 25 | train_mask = g.nodes[predict_ntype].data['train_mask'] 26 | val_mask = g.nodes[predict_ntype].data['val_mask'] 27 | test_mask = g.nodes[predict_ntype].data['test_mask'] 28 | 29 | model = HGT( 30 | {ntype: g.nodes[ntype].data['feat'].shape[1] for ntype in g.ntypes}, 31 | args.num_hidden, data.num_classes, args.num_heads, g.ntypes, g.canonical_etypes, 32 | predict_ntype, args.num_layers, args.dropout 33 | ) 34 | optimizer = optim.AdamW(model.parameters()) 35 | scheduler = optim.lr_scheduler.OneCycleLR(optimizer, args.max_lr, total_steps=args.epochs) 36 | metrics = 'Epoch {:d} | Train Loss {:.4f} | Train Micro-F1 {:.4f} | Train Macro-F1 {:.4f}' \ 37 | ' | Val Micro-F1 {:.4f} | Val Macro-F1 {:.4f}' \ 38 | ' | Test Micro-F1 {:.4f} | Test Macro-F1 {:.4f}' 39 | warnings.filterwarnings('ignore', 'Setting attributes on ParameterDict is not supported') 40 | for epoch in range(args.epochs): 41 | model.train() 42 | logits = model(g, features) 43 | loss = F.cross_entropy(logits[train_mask], labels[train_mask]) 44 | optimizer.zero_grad() 45 | loss.backward() 46 | torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) 47 | optimizer.step() 48 | scheduler.step() 49 | 50 | train_scores = micro_macro_f1_score(logits[train_mask], labels[train_mask]) 51 | val_scores = evaluate(model, g, features, labels, val_mask, micro_macro_f1_score) 52 | test_scores = evaluate(model, g, features, labels, test_mask, micro_macro_f1_score) 53 | print(metrics.format(epoch, loss.item(), *train_scores, *val_scores, *test_scores)) 54 | test_scores = evaluate(model, g, features, labels, test_mask, micro_macro_f1_score) 55 | print('Test Micro-F1 {:.4f} | Test Macro-F1 {:.4f}'.format(*test_scores)) 56 | 57 | 58 | @torch.no_grad() 59 | def evaluate(model, g, features, labels, mask, score): 60 | model.eval() 61 | logits = model(g, features) 62 | return score(logits[mask], labels[mask]) 63 | 64 | 65 | def main(): 66 | parser = argparse.ArgumentParser(description='HGT Node Classification') 67 | parser.add_argument('--seed', type=int, default=1, help='random seed') 68 | parser.add_argument('--dataset', choices=['acm', 'imdb'], default='acm', help='dataset') 69 | parser.add_argument('--num-hidden', type=int, default=256, help='number of hidden units') 70 | parser.add_argument('--num-heads', type=int, default=8, help='number of attention heads') 71 | parser.add_argument('--num-layers', type=int, default=2, help='number of layers') 72 | parser.add_argument('--dropout', type=float, default=0.5, help='dropout probability') 73 | parser.add_argument('--epochs', type=int, default=30, help='number of training epochs') 74 | parser.add_argument('--max-lr', type=float, default=1e-3, help='upper learning rate boundary') 75 | parser.add_argument('--clip', type=float, default=1.0, help='gradient norm clipping') 76 | args = parser.parse_args() 77 | print(args) 78 | train(args) 79 | 80 | 81 | if __name__ == '__main__': 82 | main() 83 | -------------------------------------------------------------------------------- /gnn/lp/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/lp/__init__.py -------------------------------------------------------------------------------- /gnn/lp/model.py: -------------------------------------------------------------------------------- 1 | """Label Propagation 2 | 3 | 论文链接：https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.14.3864&rep=rep1&type=pdf 4 | """ 5 | import dgl.function as fn 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | 11 | class LabelPropagation(nn.Module): 12 | 13 | def __init__(self, num_layers, alpha): 14 | """标签传播模型 15 | 16 | .. math:: 17 | Y^{(t+1)} = \\alpha D^{-1/2}AD^{-1/2}Y^{(t)} + (1-\\alpha)Y^{(t)} 18 | 19 | :param num_layers: int 传播层数 20 | :param alpha: float α参数 21 | """ 22 | super().__init__() 23 | self.num_layers = num_layers 24 | self.alpha = alpha 25 | 26 | @torch.no_grad() 27 | def forward(self, g, labels, mask=None): 28 | """ 29 | :param g: DGLGraph 无向图 30 | :param labels: tensor(N) 标签 31 | :param mask: tensor(N), optional 有标签顶点mask 32 | :return: tensor(N, C) 预测标签概率 33 | """ 34 | with g.local_scope(): 35 | labels = F.one_hot(labels).float() # (N, C) 36 | if mask is not None: 37 | y = torch.zeros_like(labels) 38 | y[mask] = labels[mask] 39 | else: 40 | y = labels 41 | 42 | degs = g.in_degrees().clamp(min=1) 43 | norm = torch.pow(degs, -0.5).unsqueeze(1) # D^{-1/2}, (N, 1) 44 | 45 | for _ in range(self.num_layers): 46 | g.ndata['h'] = norm * y 47 | g.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h')) 48 | y = self.alpha * norm * g.ndata.pop('h') + (1 - self.alpha) * y 49 | return y 50 | -------------------------------------------------------------------------------- /gnn/lp/readme.txt: -------------------------------------------------------------------------------- 1 | Cora 2 | python -m gnn.lp.train --dataset=cora 3 | Test Accuracy 0.6920 4 | 5 | Citeseer 6 | python -m gnn.lp.train --dataset=citeseer --num-layers=100 --alpha=0.99 7 | Test Accuracy 0.5130 8 | 9 | Pubmed 10 | python -m gnn.lp.train --dataset=pubmed --num-layers=60 --alpha=1 11 | Test Accuracy 0.7140 12 | -------------------------------------------------------------------------------- /gnn/lp/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import dgl 4 | 5 | from gnn.lp.model import LabelPropagation 6 | from gnn.utils import load_citation_dataset, accuracy 7 | 8 | 9 | def train(args): 10 | data = load_citation_dataset(args.dataset) 11 | g = data[0] 12 | labels = g.ndata['label'] 13 | train_mask = g.ndata['train_mask'] 14 | test_mask = g.ndata['test_mask'] 15 | 16 | g = dgl.add_self_loop(dgl.remove_self_loop(g)) 17 | 18 | lp = LabelPropagation(args.num_layers, args.alpha) 19 | logits = lp(g, labels, train_mask) 20 | test_acc = accuracy(logits[test_mask], labels[test_mask]) 21 | print('Test Accuracy {:.4f}'.format(test_acc)) 22 | 23 | 24 | def main(): 25 | parser = argparse.ArgumentParser(description='Label Propagation') 26 | parser.add_argument('--dataset', choices=['cora', 'citeseer', 'pubmed'], default='cora', help='dataset') 27 | parser.add_argument('--num-layers', type=int, default=10, help='number of propagation layers') 28 | parser.add_argument('--alpha', type=float, default=0.5) 29 | args = parser.parse_args() 30 | train(args) 31 | 32 | 33 | if __name__ == '__main__': 34 | main() 35 | -------------------------------------------------------------------------------- /gnn/magnn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/magnn/__init__.py -------------------------------------------------------------------------------- /gnn/magnn/encoder.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class MetapathInstanceEncoder(nn.Module): 5 | """元路径实例编码器，将一个元路径实例所有中间顶点的特征编码为一个向量。""" 6 | 7 | def forward(self, feat): 8 | """ 9 | :param feat: tensor(E, L, d_in) 10 | :return: tensor(E, d_out) 11 | """ 12 | raise NotImplementedError 13 | 14 | 15 | class MeanEncoder(MetapathInstanceEncoder): 16 | 17 | def __init__(self, in_dim, out_dim): 18 | super().__init__() 19 | 20 | def forward(self, feat): 21 | return feat.mean(dim=1) 22 | 23 | 24 | class LinearEncoder(MetapathInstanceEncoder): 25 | 26 | def __init__(self, in_dim, out_dim): 27 | super().__init__() 28 | self.fc = nn.Linear(in_dim, out_dim) 29 | 30 | def forward(self, feat): 31 | return self.fc(feat.mean(dim=1)) 32 | 33 | 34 | ENCODERS = { 35 | 'mean': MeanEncoder, 36 | 'linear': LinearEncoder 37 | } 38 | 39 | 40 | def get_encoder(name, in_dim, out_dim): 41 | if name in ENCODERS: 42 | return ENCODERS[name](in_dim, out_dim) 43 | else: 44 | raise ValueError('非法编码器名称{}，可选项为{}'.format(name, list(ENCODERS.keys()))) 45 | -------------------------------------------------------------------------------- /gnn/magnn/readme.txt: -------------------------------------------------------------------------------- 1 | 顶点分类 2 | DBLP 3 | python -m gnn.magnn.train_dblp 4 | 5 | IMDb 6 | python -m gnn.magnn.train_imdb 7 | 8 | 顶点聚类 9 | DBLP 10 | 11 | IMDb 12 | 13 | 连接预测 14 | Last.FM 15 | -------------------------------------------------------------------------------- /gnn/magnn/train_dblp.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import torch 4 | import torch.nn.functional as F 5 | import torch.optim as optim 6 | from dgl.dataloading import NodeCollator, MultiLayerNeighborSampler 7 | from torch.utils.data.dataloader import DataLoader 8 | 9 | from gnn.data import DBLPFourAreaDataset 10 | from gnn.magnn.encoder import ENCODERS 11 | from gnn.magnn.model import MAGNNMinibatch 12 | from gnn.utils import set_random_seed, metapath_based_graph, to_ntype_list, \ 13 | micro_macro_f1_score 14 | 15 | 16 | def train(args): 17 | set_random_seed(args.seed) 18 | data = DBLPFourAreaDataset() 19 | g = data[0] 20 | metapaths = data.metapaths 21 | predict_ntype = data.predict_ntype 22 | generate_one_hot_id(g) 23 | features = g.ndata['feat'] # Dict[str, tensor(N_i, d_i)] 24 | labels = g.nodes[predict_ntype].data['label'] 25 | train_idx = g.nodes[predict_ntype].data['train_mask'].nonzero(as_tuple=True)[0] 26 | val_idx = g.nodes[predict_ntype].data['val_mask'].nonzero(as_tuple=True)[0] 27 | test_idx = g.nodes[predict_ntype].data['test_mask'].nonzero(as_tuple=True)[0] 28 | out_shape = (g.num_nodes(predict_ntype), data.num_classes) 29 | 30 | print('正在生成基于元路径的图（有点慢）...') 31 | mgs = [metapath_based_graph(g, metapath) for metapath in metapaths] 32 | mgs[0].ndata['feat'] = features[predict_ntype] 33 | sampler = MultiLayerNeighborSampler([args.neighbor_size]) 34 | collators = [NodeCollator(mg, None, sampler) for mg in mgs] 35 | train_dataloader = DataLoader(train_idx, batch_size=args.batch_size) 36 | val_dataloader = DataLoader(val_idx, batch_size=args.batch_size) 37 | test_dataloader = DataLoader(test_idx, batch_size=args.batch_size) 38 | 39 | metapaths_ntype = [to_ntype_list(g, metapath) for metapath in metapaths] 40 | model = MAGNNMinibatch( 41 | predict_ntype, metapaths_ntype, 42 | {ntype: feat.shape[1] for ntype, feat in features.items()}, 43 | args.num_hidden, data.num_classes, args.num_heads, args.encoder, args.dropout 44 | ) 45 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 46 | for epoch in range(args.epochs): 47 | model.train() 48 | losses = [] 49 | train_logits = torch.zeros(out_shape) 50 | for batch in train_dataloader: 51 | gs = [collator.collate(batch)[2][0] for collator in collators] 52 | train_logits[batch] = logits = model(gs, features) 53 | loss = F.cross_entropy(logits, labels[batch]) 54 | losses.append(loss.item()) 55 | optimizer.zero_grad() 56 | loss.backward() 57 | optimizer.step() 58 | 59 | train_metrics = micro_macro_f1_score(train_logits[train_idx], labels[train_idx]) 60 | print('Epoch {:d} | Train Loss {:.4f} | Train Micro-F1 {:.4f} | Train Macro-F1 {:.4f}'.format( 61 | epoch, torch.tensor(losses).mean().item(), *train_metrics 62 | )) 63 | if (epoch + 1) % 10 == 0: 64 | val_metrics = evaluate(out_shape, collators, val_dataloader, model, features, labels) 65 | print('Val Micro-F1 {:.4f} | Val Macro-F1 {:.4f}'.format(*val_metrics)) 66 | 67 | test_metrics = evaluate(out_shape, collators, test_dataloader, model, features, labels) 68 | print('Test Micro-F1 {:.4f} | Test Macro-F1 {:.4f}'.format(*test_metrics)) 69 | 70 | 71 | def generate_one_hot_id(g): 72 | for ntype in g.ntypes: 73 | if 'feat' not in g.nodes[ntype].data: 74 | g.nodes[ntype].data['feat'] = torch.eye(g.num_nodes(ntype)) 75 | 76 | 77 | def evaluate(out_shape, collators, dataloader, model, features, labels): 78 | logits = torch.zeros(out_shape) 79 | idx = dataloader.dataset 80 | model.eval() 81 | with torch.no_grad(): 82 | for batch in dataloader: 83 | gs = [collator.collate(batch)[2][0] for collator in collators] 84 | logits[batch] = model(gs, features) 85 | return micro_macro_f1_score(logits[idx], labels[idx]) 86 | 87 | 88 | def main(): 89 | parser = argparse.ArgumentParser(description='MAGNN Node Classification (minibatch)') 90 | parser.add_argument('--seed', type=int, default=1, help='random seed') 91 | parser.add_argument('--num-hidden', type=int, default=64, help='dimension of hidden state') 92 | parser.add_argument('--num-heads', type=int, default=8, help='number of attention heads') 93 | parser.add_argument( 94 | '--encoder', choices=list(ENCODERS.keys()), default='linear', 95 | help='metapath instance encoder' 96 | ) 97 | parser.add_argument('--dropout', type=float, default=0.5, help='feature and attention dropout') 98 | parser.add_argument('--epochs', type=int, default=100, help='number of training epochs') 99 | parser.add_argument('--batch-size', type=int, default=8, help='batch size') 100 | parser.add_argument('--neighbor-size', type=int, default=100, help='neighbor sample size') 101 | parser.add_argument('--lr', type=float, default=0.005, help='learning rate') 102 | parser.add_argument('--weight-decay', type=float, default=0.001, help='weight decay') 103 | args = parser.parse_args() 104 | print(args) 105 | train(args) 106 | 107 | 108 | if __name__ == '__main__': 109 | main() 110 | -------------------------------------------------------------------------------- /gnn/magnn/train_imdb.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import torch 4 | import torch.nn.functional as F 5 | import torch.optim as optim 6 | 7 | from gnn.data import IMDbDataset 8 | from gnn.magnn.encoder import ENCODERS 9 | from gnn.magnn.model import MAGNNMultiLayer 10 | from gnn.utils import set_random_seed, metapath_based_graph, to_ntype_list, micro_macro_f1_score 11 | 12 | METAPATHS = { 13 | 'movie': [['ma', 'am'], ['md', 'dm']], 14 | 'director': [['dm', 'md'], ['dm', 'ma', 'am', 'md']], 15 | 'actor': [['am', 'ma'], ['am', 'md', 'dm', 'ma']] 16 | } 17 | 18 | 19 | def train(args): 20 | set_random_seed(args.seed) 21 | data = IMDbDataset() 22 | g = data[0] 23 | predict_ntype = data.predict_ntype 24 | features = g.ndata['feat'] # Dict[str, tensor(N_i, d_i)] 25 | labels = g.nodes[predict_ntype].data['label'] 26 | train_mask = g.nodes[predict_ntype].data['train_mask'] 27 | val_mask = g.nodes[predict_ntype].data['val_mask'] 28 | test_mask = g.nodes[predict_ntype].data['test_mask'] 29 | 30 | print('正在生成基于元路径的图...') 31 | mgs = { 32 | ntype: [metapath_based_graph(g, metapath) for metapath in METAPATHS[ntype]] 33 | for ntype in METAPATHS 34 | } 35 | for ntype in mgs: 36 | mgs[ntype][0].ndata['feat'] = g.nodes[ntype].data['feat'] 37 | metapaths_ntype = { 38 | ntype: [to_ntype_list(g, metapath) for metapath in METAPATHS[ntype]] 39 | for ntype in METAPATHS 40 | } 41 | 42 | model = MAGNNMultiLayer( 43 | args.num_layers, metapaths_ntype, 44 | {ntype: feat.shape[1] for ntype, feat in features.items()}, 45 | args.num_hidden, data.num_classes, args.num_heads, args.encoder, args.dropout 46 | ) 47 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 48 | for epoch in range(args.epochs): 49 | model.train() 50 | logits = model(mgs, features)[predict_ntype] 51 | loss = F.cross_entropy(logits[train_mask], labels[train_mask]) 52 | optimizer.zero_grad() 53 | loss.backward() 54 | optimizer.step() 55 | 56 | train_metrics = micro_macro_f1_score(logits[train_mask], labels[train_mask]) 57 | print('Epoch {:d} | Train Loss {:.4f} | Train Micro-F1 {:.4f} | Train Macro-F1 {:.4f}'.format( 58 | epoch, loss.item(), *train_metrics 59 | )) 60 | if (epoch + 1) % 10 == 0: 61 | val_metrics = evaluate(model, mgs, features, predict_ntype, labels, val_mask) 62 | print('Val Micro-F1 {:.4f} | Val Macro-F1 {:.4f}'.format(*val_metrics)) 63 | 64 | test_metrics = evaluate(model, mgs, features, predict_ntype, labels, test_mask) 65 | print('Test Micro-F1 {:.4f} | Test Macro-F1 {:.4f}'.format(*test_metrics)) 66 | 67 | 68 | def evaluate(model, gs, features, predict_ntype, labels, mask): 69 | model.eval() 70 | with torch.no_grad(): 71 | logits = model(gs, features)[predict_ntype] 72 | return micro_macro_f1_score(logits[mask], labels[mask]) 73 | 74 | 75 | def main(): 76 | parser = argparse.ArgumentParser(description='MAGNN Node Classification (multi-layer)') 77 | parser.add_argument('--seed', type=int, default=1, help='random seed') 78 | parser.add_argument('--num-layers', type=int, default=2, help='number of MAGNN layers') 79 | parser.add_argument('--num-hidden', type=int, default=64, help='dimension of hidden state') 80 | parser.add_argument('--num-heads', type=int, default=8, help='number of attention heads') 81 | parser.add_argument( 82 | '--encoder', choices=list(ENCODERS.keys()), default='linear', 83 | help='metapath instance encoder' 84 | ) 85 | parser.add_argument('--dropout', type=float, default=0.5, help='feature and attention dropout') 86 | parser.add_argument('--epochs', type=int, default=100, help='number of training epochs') 87 | parser.add_argument('--lr', type=float, default=0.005, help='learning rate') 88 | parser.add_argument('--weight-decay', type=float, default=0.001, help='weight decay') 89 | args = parser.parse_args() 90 | print(args) 91 | train(args) 92 | 93 | 94 | if __name__ == '__main__': 95 | main() 96 | -------------------------------------------------------------------------------- /gnn/metapath2vec/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/metapath2vec/__init__.py -------------------------------------------------------------------------------- /gnn/metapath2vec/random_walk.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import dgl 4 | import torch 5 | from dgl.sampling import random_walk 6 | from tqdm import trange 7 | 8 | from gnn.data import AMinerCSDataset 9 | from gnn.utils import metapath_adj 10 | 11 | 12 | def main(): 13 | parser = argparse.ArgumentParser(description='Metapath-based Random Walk for metapath2vec') 14 | parser.add_argument('--num-walks', type=int, default=1000, help='number of walks for each node') 15 | parser.add_argument('--walk-length', type=int, default=100, help='times to repeat metapath') 16 | parser.add_argument('output_file', help='output filename') 17 | args = parser.parse_args() 18 | 19 | data = AMinerCSDataset() 20 | g = data[0] 21 | 22 | ca = metapath_adj(g, ['cp', 'pa']) 23 | ca_c, ca_a = ca.nonzero() 24 | cag = dgl.heterograph( 25 | {('conf', 'ca', 'author'): (ca_c, ca_a), ('author', 'ac', 'conf'): (ca_a, ca_c)}, 26 | {'conf': ca.shape[0], 'author': ca.shape[1]} 27 | ) 28 | cag.edges['ca'].data['p'] = cag.edges['ac'].data['p'] = torch.from_numpy(ca.data).float() 29 | 30 | metapath = ['ca', 'ac'] # metapath = CAC, metapath*2 = CACAC 31 | f = open(args.output_file, 'w') 32 | for cid in trange(cag.num_nodes('conf'), ncols=80): 33 | traces, _ = random_walk( 34 | cag, [cid] * args.num_walks, metapath=metapath * args.walk_length, prob='p' 35 | ) 36 | f.writelines([trace2name(data.author_names, data.conf_names, t) + '\n' for t in traces]) 37 | f.close() 38 | 39 | 40 | def trace2name(author_names, conf_names, trace): 41 | return ' '.join((author_names if i & 1 else conf_names)[trace[i]] for i in range(len(trace))) 42 | 43 | 44 | if __name__ == '__main__': 45 | main() 46 | -------------------------------------------------------------------------------- /gnn/metapath2vec/readme.txt: -------------------------------------------------------------------------------- 1 | 1.随机游走生成语料库 2 | python -m gnn.metapath2vec.random_walk --num-walks=1000 --walk-length=100 D:\dgldata\aminer-cs\corpus.txt 3 | 4 | 2.训练word2vec 5 | python -m gnn.metapath2vec.train_word2vec --size=128 --workers=8 D:\dgldata\aminer-cs\corpus.txt D:\dgldata\aminer-cs\aminer-cs.model 6 | 7 | 3.评价顶点嵌入 8 | 期刊（会议）顶点分类 9 | python -m gnn.metapath2vec.train --ntype=conf --task=clf D:\dgldata\aminer-cs\aminer-cs.model 10 | Train size 5% | Test Micro-F1 0.4453 | Test Macro-F1 0.3634 11 | Train size 10% | Test Micro-F1 0.5620 | Test Macro-F1 0.4602 12 | Train size 20% | Test Micro-F1 0.6667 | Test Macro-F1 0.6281 13 | Train size 30% | Test Micro-F1 0.9043 | Test Macro-F1 0.9096 14 | Train size 40% | Test Micro-F1 0.9259 | Test Macro-F1 0.9273 15 | Train size 50% | Test Micro-F1 0.9552 | Test Macro-F1 0.9563 16 | Train size 60% | Test Micro-F1 0.9815 | Test Macro-F1 0.9838 17 | Train size 70% | Test Micro-F1 1.0000 | Test Macro-F1 1.0000 18 | Train size 80% | Test Micro-F1 1.0000 | Test Macro-F1 1.0000 19 | Train size 90% | Test Micro-F1 1.0000 | Test Macro-F1 1.0000 20 | 21 | 学者顶点分类 22 | python -m gnn.metapath2vec.train --ntype=author --task=clf D:\dgldata\aminer-cs\aminer-cs.model 23 | Train size 5% | Test Micro-F1 0.9334 | Test Macro-F1 0.9274 24 | Train size 10% | Test Micro-F1 0.9395 | Test Macro-F1 0.9343 25 | Train size 20% | Test Micro-F1 0.9418 | Test Macro-F1 0.9368 26 | Train size 30% | Test Micro-F1 0.9426 | Test Macro-F1 0.9380 27 | Train size 40% | Test Micro-F1 0.9429 | Test Macro-F1 0.9385 28 | Train size 50% | Test Micro-F1 0.9433 | Test Macro-F1 0.9389 29 | Train size 60% | Test Micro-F1 0.9433 | Test Macro-F1 0.9387 30 | Train size 70% | Test Micro-F1 0.9425 | Test Macro-F1 0.9381 31 | Train size 80% | Test Micro-F1 0.9427 | Test Macro-F1 0.9385 32 | Train size 90% | Test Micro-F1 0.9414 | Test Macro-F1 0.9371 33 | 34 | 期刊（会议）顶点聚类 35 | python -m gnn.metapath2vec.train --ntype=conf --task=cluster D:\dgldata\aminer-cs\aminer-cs.model 36 | Average NMI 0.9173 37 | 38 | 学者顶点聚类 39 | python -m gnn.metapath2vec.train --ntype=author --task=cluster D:\dgldata\aminer-cs\aminer-cs.model 40 | Average NMI 0.7497 41 | 42 | 43 | 注意：原论文提供的随机游走代码py4genMetaPaths.py中预先计算出学者-期刊对应关系，每一步随机游走是从对应关系中等概率选择 44 | 这与dgl.sampling.random_walk()选择顶点的概率并不相等！ 45 | 46 | 例如，有以下异构图 47 | g = dgl.heterograph({ 48 | ('author', 'ap', 'paper'): ([0, 0, 1, 1, 2], [0, 1, 1, 2, 2]), 49 | ('paper', 'pa', 'author'): ([0, 1, 1, 2, 2], [0, 0, 1, 1, 2]), 50 | ('paper', 'pc', 'conf'): ([0, 1, 2], [0, 0, 1]), 51 | ('conf', 'cp', 'paper'): ([0, 0, 1], [0, 1, 2]) 52 | }) 53 | 54 | a0 - p0 55 | \ \ 56 | a1 - p1 - c0 57 | \ 58 | a2 - p2 - c1 59 | 60 | （1）原论文代码：计算学者-期刊对应关系 61 | conf_author = {c0: [a0, a0, a1], c1: [a1, a2]} 62 | author_conf = {a0: [c0, c0], a1: [c0, c1], a2: [c1]} 63 | c0 -> 2/3a0 + 1/3a1 64 | 65 | （2）dgl.sampling.random_walk(g, [0], metapath=['cp', 'pa']) 66 | c0 -> 1/2p0 + 1/2p1 -> 1/2a0 + 1/4a0 + 1/4a1 = 3/4a0 + 1/4a1 67 | 68 | >>> traces, _ = dgl.sampling.random_walk(g, [0] * 1000, metapath=['cp', 'pa']) 69 | >>> authors = traces[:, 2] 70 | >>> authors.unique(return_counts=True) 71 | (tensor([0, 1]), tensor([754, 246])) 72 | 73 | 解决方法：使用prob参数 74 | >>> cp = g.adj(etype='cp', scipy_fmt='csr') 75 | >>> pa = g.adj(etype='pa', scipy_fmt='csr') 76 | >>> ca = cp * pa 77 | >>> ca.todense() 78 | matrix([[2, 1, 0], 79 | [0, 1, 1]], dtype=int64) 80 | >>> cag = dgl.heterograph({('conf', 'ca', 'author'): ca.nonzero()}) 81 | >>> cag.edata['p'] = torch.from_numpy(ca.data).float() 82 | >>> traces, _ = dgl.sampling.random_walk(cag, [0] * 1000, metapath=['ca'], prob='p') 83 | >>> traces[:, 1].unique(return_counts=True) 84 | (tensor([0, 1]), tensor([670, 330])) 85 | -------------------------------------------------------------------------------- /gnn/metapath2vec/skipgram.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class SkipGram(nn.Module): 7 | 8 | def __init__(self, vocab_size, embed_dim): 9 | super().__init__() 10 | self.center_embed = nn.Embedding(vocab_size, embed_dim) 11 | self.neigh_embed = nn.Embedding(vocab_size, embed_dim) 12 | 13 | def forward(self, center, pos, neg): 14 | r"""给定中心词、正样本和负样本，返回似然函数的相反数（损失）： 15 | 16 | .. math:: 17 | L=-\log {\sigma(v_c \cdot v_p)}-\sum_{n \in neg}{\log {\sigma(-v_c \cdot v_n)}} 18 | 19 | :param center: tensor(N) 中心词 20 | :param pos: tensor(N) 正样本 21 | :param neg: tensor(N, M) 负样本 22 | """ 23 | center_embed = self.center_embed(center) # (N, d) 24 | pos_embed = self.neigh_embed(pos) # (N, d) 25 | neg_embed = self.neigh_embed(neg) # (N, M, d) 26 | 27 | pos_score = torch.sum(center_embed * pos_embed, dim=1) # (N) 28 | pos_score = F.logsigmoid(torch.clamp(pos_score, min=-10, max=10)) 29 | 30 | neg_score = torch.bmm(neg_embed, center_embed.unsqueeze(2)).squeeze() # (N, M) 31 | neg_score = torch.sum(F.logsigmoid(torch.clamp(neg_score, min=-10, max=10)), dim=1) # (N) 32 | return -torch.mean(pos_score + neg_score) 33 | -------------------------------------------------------------------------------- /gnn/metapath2vec/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import random 3 | 4 | import torch 5 | from gensim.models import Word2Vec 6 | from sklearn.cluster import KMeans 7 | from sklearn.linear_model import LogisticRegression 8 | from sklearn.metrics import f1_score, normalized_mutual_info_score 9 | from sklearn.model_selection import train_test_split 10 | 11 | from gnn.data import AMinerCSDataset 12 | 13 | 14 | def load_data(ntype, model_path): 15 | data = AMinerCSDataset() 16 | g = data[0] 17 | idx = torch.nonzero(g.nodes[ntype].data['label'] != -1, as_tuple=True)[0] 18 | labels = g.nodes[ntype].data['label'][idx].numpy() 19 | 20 | model = Word2Vec.load(model_path) 21 | names = data.author_names if ntype == 'author' else data.conf_names 22 | embeds = model.wv[[names[i] for i in idx]] 23 | return embeds, labels, data.num_classes 24 | 25 | 26 | def node_clf(args): 27 | embeds, labels, _ = load_data(args.ntype, args.model_path) 28 | for train_size in [0.05] + [0.1 * x for x in range(1, 10)]: 29 | X_train, X_test, y_train, y_test = train_test_split( 30 | embeds, labels, train_size=train_size, random_state=args.seed 31 | ) 32 | clf = LogisticRegression(random_state=args.seed, max_iter=500) 33 | clf.fit(X_train, y_train) 34 | y_pred = clf.predict(X_test) 35 | micro_f1 = f1_score(y_test, y_pred, average='micro') 36 | macro_f1 = f1_score(y_test, y_pred, average='macro') 37 | print('Train size {:.0%} | Test Micro-F1 {:.4f} | Test Macro-F1 {:.4f}'.format( 38 | train_size, micro_f1, macro_f1 39 | )) 40 | 41 | 42 | def node_cluster(args): 43 | embeds, labels, num_classes = load_data(args.ntype, args.model_path) 44 | random.seed(args.seed) 45 | seeds = [random.randint(0, 0x7fffffff) for _ in range(10)] 46 | nmis = [] 47 | for seed in seeds: 48 | pred = KMeans(num_classes, random_state=seed).fit_predict(embeds) 49 | nmi = normalized_mutual_info_score(labels, pred) 50 | print('Seed {:d} | NMI {:.4f}'.format(seed, nmi)) 51 | nmis.append(nmi) 52 | print('Average NMI {:.4f}'.format(sum(nmis) / len(nmis))) 53 | 54 | 55 | def main(): 56 | parser = argparse.ArgumentParser(description='metapath2vec Node Classification or Clustering') 57 | parser.add_argument('--seed', type=int, default=2, help='random seed') 58 | parser.add_argument('--ntype', choices=['author', 'conf'], default='author', help='node type') 59 | parser.add_argument('--task', choices=['clf', 'cluster'], default='clf', help='training task') 60 | parser.add_argument('model_path', help='path to word2vec model (node embeddings)') 61 | args = parser.parse_args() 62 | if args.task == 'clf': 63 | node_clf(args) 64 | else: 65 | node_cluster(args) 66 | 67 | 68 | if __name__ == '__main__': 69 | main() 70 | -------------------------------------------------------------------------------- /gnn/metapath2vec/train_word2vec.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | 4 | from gensim.models import Word2Vec 5 | 6 | logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO) 7 | 8 | 9 | def main(): 10 | parser = argparse.ArgumentParser(description='Train word2vec for metapath2vec') 11 | parser.add_argument('--size', type=int, default=128, help='dimensionality of the word vectors') 12 | parser.add_argument('--workers', type=int, default=3, help='number of worker threads') 13 | parser.add_argument('--iter', type=int, default=10, help='number of iterations') 14 | parser.add_argument('corpus_file', help='path to corpus file') 15 | parser.add_argument('save_path', help='path to file where to save word2vec model') 16 | args = parser.parse_args() 17 | print(args) 18 | 19 | model = Word2Vec( 20 | corpus_file=args.corpus_file, size=args.size, min_count=1, 21 | workers=args.workers, sg=1, iter=args.iter 22 | ) 23 | model.save(args.save_path) 24 | 25 | 26 | if __name__ == '__main__': 27 | main() 28 | -------------------------------------------------------------------------------- /gnn/rgcn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/rgcn/__init__.py -------------------------------------------------------------------------------- /gnn/rgcn/model.py: -------------------------------------------------------------------------------- 1 | """Relational Graph Convolutional Network (R-GCN) 2 | 3 | 论文链接：https://arxiv.org/abs/1703.06103 4 | """ 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from dgl.nn import RelGraphConv 9 | 10 | 11 | class DistMult(nn.Module): 12 | 13 | def __init__(self, num_rels, feat_dim): 14 | """知识图谱嵌入模型DistMult 15 | 16 | :param num_rels: int 关系个数 17 | :param feat_dim: int 嵌入维数 18 | """ 19 | super().__init__() 20 | self.w_relations = nn.Parameter(torch.Tensor(num_rels, feat_dim)) 21 | self.reset_parameters() 22 | 23 | def reset_parameters(self): 24 | nn.init.xavier_uniform_(self.w_relations, gain=nn.init.calculate_gain('relu')) 25 | 26 | def forward(self, embed, head, rel, tail): 27 | """ 28 | :param embed: tensor(N, d) 实体嵌入 29 | :param head: tensor(*) 头实体 30 | :param rel: tensor(*) 关系 31 | :param tail: tensor(*) 尾实体 32 | :return: tensor(*) 三元组得分 33 | """ 34 | # e_h^T R_r e_t = (e_h * diag(R_r)) e_t 35 | return torch.sum(embed[head] * self.w_relations[rel] * embed[tail], dim=1) 36 | 37 | 38 | class LinkPrediction(nn.Module): 39 | 40 | def __init__( 41 | self, num_nodes, hidden_dim, num_rels, num_layers=2, 42 | regularizer='basis', num_bases=None, dropout=0.0): 43 | """R-GCN连接预测模型 44 | 45 | :param num_nodes: int 顶点（实体）数 46 | :param hidden_dim: int 隐含特征维数 47 | :param num_rels: int 关系个数 48 | :param num_layers: int, optional R-GCN层数，默认为2 49 | :param regularizer: str, 'basis'/'bdd' 权重正则化方法，默认为'basis' 50 | :param num_bases: int, optional 基的个数，默认使用关系个数 51 | :param dropout: float, optional Dropout概率，默认为0 52 | """ 53 | super().__init__() 54 | self.embed = nn.Embedding(num_nodes, hidden_dim) 55 | self.layers = nn.ModuleList() 56 | for i in range(num_layers): 57 | self.layers.append(RelGraphConv( 58 | hidden_dim, hidden_dim, num_rels, regularizer, num_bases, 59 | activation=F.relu if i < num_layers - 1 else None, 60 | self_loop=True, low_mem=True, dropout=dropout 61 | )) 62 | self.score = DistMult(num_rels, hidden_dim) 63 | 64 | def forward(self, g, etypes): 65 | """ 66 | :param g: DGLGraph 同构图 67 | :param etypes: tensor(|E|) 边类型 68 | :return: tensor(N, d_hid) 顶点嵌入 69 | """ 70 | h = self.embed.weight # (N, d_hid) 71 | for layer in self.layers: 72 | h = layer(g, h, etypes) 73 | return h 74 | 75 | def calc_score(self, embed, triplets): 76 | """计算三元组得分 77 | 78 | :param embed: tensor(N, d_hid) 顶点（实体）嵌入 79 | :param triplets: (tensor(*), tensor(*), tensor(*)) 三元组(head, tail, relation) 80 | :return: tensor(*) 三元组得分 81 | """ 82 | # 三元组的relation放在最后是为了与DGLGraph.find_edges的返回格式保持一致 83 | head, tail, rel = triplets 84 | return self.score(embed, head, rel, tail) 85 | -------------------------------------------------------------------------------- /gnn/rgcn/model_hetero.py: -------------------------------------------------------------------------------- 1 | """Relational Graph Convolutional Network (R-GCN) 2 | 3 | 论文链接：https://arxiv.org/abs/1703.06103 4 | """ 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from dgl.nn import HeteroGraphConv, GraphConv, WeightBasis 9 | 10 | 11 | class RelGraphConvHetero(nn.Module): 12 | 13 | def __init__( 14 | self, in_dim, out_dim, rel_names, num_bases=None, 15 | weight=True, self_loop=True, activation=None, dropout=0.0): 16 | """R-GCN层（用于异构图） 17 | 18 | :param in_dim: 输入特征维数 19 | :param out_dim: 输出特征维数 20 | :param rel_names: List[str] 关系名称 21 | :param num_bases: int, optional 基的个数，默认使用关系个数 22 | :param weight: bool, optional 是否进行线性变换，默认为True 23 | :param self_loop: 是否包括自环消息，默认为True 24 | :param activation: callable, optional 激活函数，默认为None 25 | :param dropout: float, optional Dropout概率，默认为0 26 | """ 27 | super().__init__() 28 | self.rel_names = rel_names 29 | self.self_loop = self_loop 30 | self.activation = activation 31 | self.dropout = nn.Dropout(dropout) 32 | 33 | self.conv = HeteroGraphConv({ 34 | rel: GraphConv(in_dim, out_dim, norm='right', weight=False, bias=False) 35 | for rel in rel_names 36 | }) 37 | 38 | self.use_weight = weight 39 | if not num_bases: 40 | num_bases = len(rel_names) 41 | self.use_basis = weight and 0 < num_bases < len(rel_names) 42 | if self.use_weight: 43 | if self.use_basis: 44 | self.basis = WeightBasis((in_dim, out_dim), num_bases, len(rel_names)) 45 | else: 46 | self.weight = nn.Parameter(torch.Tensor(len(rel_names), in_dim, out_dim)) 47 | nn.init.xavier_uniform_(self.weight, nn.init.calculate_gain('relu')) 48 | 49 | if self.self_loop: 50 | self.loop_weight = nn.Parameter(torch.Tensor(in_dim, out_dim)) 51 | nn.init.xavier_uniform_(self.loop_weight, nn.init.calculate_gain('relu')) 52 | 53 | def forward(self, g, inputs): 54 | """ 55 | :param g: DGLGraph 异构图 56 | :param inputs: Dict[str, tensor(N_i, d_in)] 顶点类型到输入特征的映射 57 | :return: Dict[str, tensor(N_i, d_out)] 顶点类型到输出特征的映射 58 | """ 59 | if self.use_weight: 60 | weight = self.basis() if self.use_basis else self.weight # (R, d_in, d_out) 61 | kwargs = {rel: {'weight': weight[i]} for i, rel in enumerate(self.rel_names)} 62 | else: 63 | kwargs = {} 64 | hs = self.conv(g, inputs, mod_kwargs=kwargs) # Dict[ntype, (N_i, d_out)] 65 | for ntype in hs: 66 | if self.self_loop: 67 | hs[ntype] += torch.matmul(inputs[ntype], self.loop_weight) 68 | if self.activation: 69 | hs[ntype] = self.activation(hs[ntype]) 70 | hs[ntype] = self.dropout(hs[ntype]) 71 | return hs 72 | 73 | 74 | class EntityClassification(nn.Module): 75 | 76 | def __init__( 77 | self, num_nodes, hidden_dim, out_dim, rel_names, 78 | num_hidden_layers=1, num_bases=None, self_loop=True, dropout=0.0): 79 | """R-GCN实体分类模型 80 | 81 | :param num_nodes: Dict[str, int] 顶点类型到顶点数的映射 82 | :param hidden_dim: int 隐含特征维数 83 | :param out_dim: int 输出特征维数 84 | :param rel_names: List[str] 关系名称 85 | :param num_hidden_layers: int, optional R-GCN隐藏层数，默认为1 86 | :param num_bases: int, optional 基的个数，默认使用关系个数 87 | :param self_loop: bool 是否包括自环消息，默认为True 88 | :param dropout: float, optional Dropout概率，默认为0 89 | """ 90 | super().__init__() 91 | self.embeds = nn.ModuleDict({ 92 | ntype: nn.Embedding(num_nodes[ntype], hidden_dim) for ntype in num_nodes 93 | }) 94 | self.layers = nn.ModuleList() 95 | for _ in range(num_hidden_layers): 96 | self.layers.append(RelGraphConvHetero( 97 | hidden_dim, hidden_dim, rel_names, num_bases, False, self_loop, F.relu, dropout 98 | )) 99 | self.layers.append(RelGraphConvHetero( 100 | hidden_dim, out_dim, rel_names, num_bases, True, self_loop, dropout=dropout 101 | )) 102 | self.reset_parameters() 103 | 104 | def reset_parameters(self): 105 | # normal_ -> xavier_uniform_ => acc 0.8 -> 0.97? 106 | for k in self.embeds: 107 | nn.init.xavier_uniform_(self.embeds[k].weight, gain=nn.init.calculate_gain('relu')) 108 | 109 | def forward(self, g): 110 | """ 111 | :param g: DGLGraph 异构图 112 | :return: Dict[str, tensor(N_i, d_out)] 顶点类型到顶点嵌入的映射 113 | """ 114 | h = {k: self.embeds[k].weight for k in self.embeds} 115 | for layer in self.layers: 116 | h = layer(g, h) # Dict[ntype, (N_i, d_hid)] 117 | return h 118 | -------------------------------------------------------------------------------- /gnn/rgcn/readme.txt: -------------------------------------------------------------------------------- 1 | 实体分类 2 | AIFB 3 | python -m gnn.rgcn.train_entity_clf --dataset=aifb --num-hidden=16 --num-bases=0 --weight-decay=0 4 | Test Accuracy 0.9167 5 | 6 | MUTAG 7 | python -m gnn.rgcn.train_entity_clf --dataset=mutag --num-hidden=16 --num-bases=30 8 | Test Accuracy 0.7647 9 | 10 | BGS 11 | python -m gnn.rgcn.train_entity_clf --dataset=bgs --num-hidden=16 --num-bases=40 12 | Test Accuracy 0.8966 13 | 14 | AM 15 | python -m gnn.rgcn.train_entity_clf --dataset=am --num-hidden=10 --num-bases=40 16 | Test Accuracy 0.8434 17 | 18 | 连接预测 19 | WN18 20 | python -m gnn.rgcn.train_link_pred --dataset=wn18 --num-hidden=200 --num-layers=1 --num-bases=2 --dropout=0.4 21 | 22 | FB15k 23 | python -m gnn.rgcn.train_link_pred --dataset=FB15k --num-hidden=200 --num-layers=1 --num-bases=2 --dropout=0.4 24 | 25 | FB15k-237 26 | python -m gnn.rgcn.train_link_pred --dataset=FB15k-237 --num-hidden=500 --num-layers=2 --regularizer=bdd --num-bases=100 --dropout=0.4 27 | -------------------------------------------------------------------------------- /gnn/rgcn/train_entity_clf.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import torch.nn.functional as F 4 | import torch.optim as optim 5 | 6 | from gnn.rgcn.model_hetero import EntityClassification 7 | from gnn.utils import load_rdf_dataset, accuracy 8 | 9 | 10 | def train(args): 11 | data = load_rdf_dataset(args.dataset) 12 | g = data[0] 13 | category = data.predict_category 14 | num_classes = data.num_classes 15 | labels = g.nodes[category].data['labels'] 16 | train_mask = g.nodes[category].data['train_mask'].bool() 17 | test_mask = g.nodes[category].data['test_mask'].bool() 18 | 19 | model = EntityClassification( 20 | {ntype: g.num_nodes(ntype) for ntype in g.ntypes}, 21 | args.num_hidden, num_classes, list(set(g.etypes)), 22 | args.num_hidden_layers, args.num_bases, args.self_loop, args.dropout 23 | ) 24 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 25 | for epoch in range(args.epochs): 26 | model.train() 27 | logits = model(g)[category] 28 | loss = F.cross_entropy(logits[train_mask], labels[train_mask]) 29 | optimizer.zero_grad() 30 | loss.backward() 31 | optimizer.step() 32 | 33 | train_acc = accuracy(logits[train_mask], labels[train_mask]) 34 | print('Epoch {:04d} | Loss {:.4f} | Train Acc {:.4f}'.format(epoch, loss.item(), train_acc)) 35 | 36 | test_acc = accuracy(model(g)[category][test_mask], labels[test_mask]) 37 | print('Test Accuracy {:.4f}'.format(test_acc)) 38 | 39 | 40 | def main(): 41 | parser = argparse.ArgumentParser(description='R-GCN Entity Classification') 42 | parser.add_argument( 43 | '--dataset', choices=['aifb', 'mutag', 'bgs', 'am'], default='aifb', help='dataset' 44 | ) 45 | parser.add_argument('--num-hidden', type=int, default=16, help='number of hidden units') 46 | parser.add_argument('--num-hidden-layers', type=int, default=1, help='number of hidden layers') 47 | parser.add_argument('--num-bases', type=int, default=0) 48 | parser.add_argument('--self-loop', action='store_true', help='include self-loop message') 49 | parser.add_argument('--dropout', type=float, default=0.0, help='dropout probability') 50 | parser.add_argument('--epochs', type=int, default=50, help='number of training epochs') 51 | parser.add_argument('--lr', type=float, default=0.01, help='learning rate') 52 | parser.add_argument('--weight-decay', type=float, default=5e-4, help='weight decay') 53 | args = parser.parse_args() 54 | print(args) 55 | train(args) 56 | 57 | 58 | if __name__ == '__main__': 59 | main() 60 | -------------------------------------------------------------------------------- /gnn/rgcn/train_link_pred.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import dgl 4 | import torch 5 | import torch.nn.functional as F 6 | import torch.optim as optim 7 | from dgl.dataloading.negative_sampler import Uniform 8 | 9 | from gnn.rgcn.model import LinkPrediction 10 | from gnn.utils import load_kg_dataset 11 | 12 | 13 | def train(args): 14 | data = load_kg_dataset(args.dataset) 15 | g = data[0] 16 | train_idx = g.edata['train_mask'].nonzero(as_tuple=False).squeeze() 17 | val_idx = g.edata['val_mask'].nonzero(as_tuple=False).squeeze() 18 | test_idx = g.edata['test_mask'].nonzero(as_tuple=False).squeeze() 19 | 20 | train_g = dgl.edge_subgraph(g, train_idx, preserve_nodes=True) 21 | train_triplets = g.find_edges(train_idx) + (train_g.edata['etype'],) 22 | model = LinkPrediction( 23 | data.num_nodes, args.num_hidden, data.num_rels * 2, args.num_layers, 24 | args.regularizer, args.num_bases, args.dropout 25 | ) 26 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 27 | neg_sampler = Uniform(args.num_neg_samples) 28 | labels = torch.cat([torch.ones(train_g.num_edges()), torch.zeros(train_g.num_edges() * args.num_neg_samples)] ) 29 | for epoch in range(args.epochs): 30 | model.train() 31 | embed = model(train_g, train_g.edata['etype']) 32 | 33 | neg_triplets = neg_sampler(train_g, torch.arange(train_g.num_edges())) \ 34 | + (train_g.edata['etype'].repeat_interleave(args.num_neg_samples),) 35 | pos_score = model.calc_score(embed, train_triplets) 36 | neg_score = model.calc_score(embed, neg_triplets) 37 | loss = F.binary_cross_entropy_with_logits(torch.cat([pos_score, neg_score]), labels) 38 | optimizer.zero_grad() 39 | loss.backward() 40 | optimizer.step() 41 | 42 | # TODO 计算MRR 43 | # FB-15k和FB15k-237反向传播很慢？ 44 | print('Epoch {:04d} | Loss {:.4f}'.format(epoch, loss.item())) 45 | 46 | 47 | def main(): 48 | parser = argparse.ArgumentParser(description='R-GCN Link Prediction') 49 | parser.add_argument( 50 | '--dataset', choices=['wn18', 'FB15k', 'FB15k-237'], default='wn18', help='dataset' 51 | ) 52 | parser.add_argument('--num-hidden', type=int, default=200, help='number of hidden units') 53 | parser.add_argument('--num-layers', type=int, default=2, help='number of encoder layers') 54 | parser.add_argument( 55 | '--regularizer', choices=['basis', 'bdd'], default='basis', help='weight regularizer' 56 | ) 57 | parser.add_argument('--num-bases', type=int, default=0) 58 | parser.add_argument('--dropout', type=float, default=0.0, help='dropout probability') 59 | parser.add_argument('--epochs', type=int, default=50, help='number of training epochs') 60 | parser.add_argument('--num-neg-samples', type=int, default=1, help='number of negative samples') 61 | parser.add_argument('--lr', type=float, default=0.01, help='learning rate') 62 | parser.add_argument('--weight-decay', type=float, default=0.01, help='weight decay') 63 | args = parser.parse_args() 64 | print(args) 65 | train(args) 66 | 67 | 68 | if __name__ == '__main__': 69 | main() 70 | -------------------------------------------------------------------------------- /gnn/rhgnn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/rhgnn/__init__.py -------------------------------------------------------------------------------- /gnn/rhgnn/readme.txt: -------------------------------------------------------------------------------- 1 | ogbn-mag数据集 2 | 3 | 1.预训练顶点嵌入 4 | 见https://github.com/ZZy979/GNN-Recommendation/blob/main/gnnrec/hge/readme.md#%E9%A2%84%E8%AE%AD%E7%BB%83%E9%A1%B6%E7%82%B9%E5%B5%8C%E5%85%A5 5 | 6 | 2.训练模型 7 | python -m gnn.rhgnn.train /home/zzy/ogb/ /home/zzy/GNN-Recommendation/model/word2vec/ogbn-mag.model 8 | Test Acc 0.5201 9 | -------------------------------------------------------------------------------- /gnn/sign/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/sign/__init__.py -------------------------------------------------------------------------------- /gnn/sign/model.py: -------------------------------------------------------------------------------- 1 | """SIGN: Scalable Inception Graph Neural Networks (SIGN) 2 | 3 | 论文链接：https://arxiv.org/pdf/2004.11198 4 | """ 5 | import torch 6 | import torch.nn as nn 7 | 8 | 9 | class FeedForwardNet(nn.Module): 10 | """L层全连接网络""" 11 | 12 | def __init__(self, in_dim, hidden_dim, out_dim, num_layers, dropout=0.0): 13 | super().__init__() 14 | self.fc = nn.ModuleList() 15 | if num_layers == 1: 16 | self.fc.append(nn.Linear(in_dim, out_dim, bias=False)) 17 | else: 18 | self.fc.append(nn.Linear(in_dim, hidden_dim, bias=False)) 19 | for _ in range(num_layers - 2): 20 | self.fc.append(nn.Linear(hidden_dim, hidden_dim, bias=False)) 21 | self.fc.append(nn.Linear(hidden_dim, out_dim, bias=False)) 22 | self.prelu = nn.PReLU() 23 | self.dropout = nn.Dropout(dropout) 24 | 25 | def forward(self, x): 26 | """ 27 | :param x: tensor(N, d_in) 28 | :return: tensor(N, d_out) 29 | """ 30 | for i, fc in enumerate(self.fc): 31 | x = fc(x) 32 | if i < len(self.fc) - 1: 33 | x = self.dropout(self.prelu(x)) 34 | return x 35 | 36 | 37 | class SIGN(nn.Module): 38 | 39 | def __init__(self, in_dim, hidden_dim, out_dim, num_hops, num_layers, dropout=0.0): 40 | """SIGN模型 41 | 42 | :param in_dim: int 输入特征维数 43 | :param hidden_dim: int 隐含特征维数 44 | :param out_dim: int 输出特征维数 45 | :param num_hops: int 跳数r 46 | :param num_layers: int 全连接网络层数 47 | :param dropout: float, optional Dropout概率，默认为0 48 | """ 49 | super().__init__() 50 | self.inception_ffs = nn.ModuleList([ 51 | FeedForwardNet(in_dim, hidden_dim, hidden_dim, num_layers, dropout) # (4)式中的Θ_i 52 | for _ in range(num_hops + 1) 53 | ]) 54 | self.project = FeedForwardNet( 55 | (num_hops + 1) * hidden_dim, hidden_dim, out_dim, num_layers, dropout 56 | ) # (4)式中的Ω 57 | self.prelu = nn.PReLU() 58 | self.dropout = nn.Dropout(dropout) 59 | 60 | def forward(self, feats): 61 | """ 62 | :param feats: List[tensor(N, d_in)] 每一跳的邻居聚集特征，长度为r+1 63 | :return: tensor(N, d_out) 输出顶点特征 64 | """ 65 | # (N, (r+1)*d_hid) 66 | h = torch.cat([ff(feat) for ff, feat in zip(self.inception_ffs, feats)], dim=-1) 67 | out = self.project(self.dropout(self.prelu(h))) # (N, d_out) 68 | return out 69 | -------------------------------------------------------------------------------- /gnn/sign/readme.txt: -------------------------------------------------------------------------------- 1 | Cora 2 | python -m gnn.sign.train --dataset=cora --epochs=20 3 | Test Acc 0.7850 4 | 5 | Reddit 6 | python -m gnn.sign.train --dataset=reddit 7 | Test Acc 0.9527 8 | 9 | ogbn-products 10 | python -m gnn.sign.train --dataset=amazon 11 | Test Acc 0.7049 12 | -------------------------------------------------------------------------------- /gnn/sign/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import dgl.function as fn 4 | import torch 5 | import torch.nn.functional as F 6 | import torch.optim as optim 7 | from dgl.data import CoraGraphDataset, RedditDataset 8 | from ogb.nodeproppred import DglNodePropPredDataset 9 | 10 | from gnn.sign.model import SIGN 11 | from gnn.utils import set_random_seed, accuracy 12 | 13 | 14 | def load_data(dataset, ogb_root): 15 | if dataset in ('cora', 'reddit'): 16 | data = CoraGraphDataset() if dataset == 'cora' else RedditDataset(self_loop=True) 17 | g = data[0] 18 | train_idx = g.ndata['train_mask'].nonzero(as_tuple=True)[0] 19 | val_idx = g.ndata['val_mask'].nonzero(as_tuple=True)[0] 20 | test_idx = g.ndata['test_mask'].nonzero(as_tuple=True)[0] 21 | return g, g.ndata['label'], data.num_classes, train_idx, val_idx, test_idx 22 | else: 23 | data = DglNodePropPredDataset('ogbn-products', ogb_root) 24 | g, labels = data[0] 25 | split_idx = data.get_idx_split() 26 | return g, labels.squeeze(dim=-1), data.num_classes, \ 27 | split_idx['train'], split_idx['valid'], split_idx['test'] 28 | 29 | 30 | def calc_weight(g): 31 | """计算行归一化的D^(-1/2)AD(-1/2)""" 32 | with g.local_scope(): 33 | g.ndata['in_degree'] = g.in_degrees().float().pow(-0.5) 34 | g.ndata['out_degree'] = g.out_degrees().float().pow(-0.5) 35 | g.apply_edges(fn.u_mul_v('out_degree', 'in_degree', 'weight')) 36 | g.update_all(fn.copy_e('weight', 'msg'), fn.sum('msg', 'norm')) 37 | g.apply_edges(fn.e_div_v('weight', 'norm', 'weight')) 38 | return g.edata['weight'] 39 | 40 | 41 | def preprocess(g, features, num_hops): 42 | """预先计算r跳的邻居聚集特征""" 43 | with torch.no_grad(): 44 | g.edata['weight'] = calc_weight(g) 45 | g.ndata['feat_0'] = features 46 | for h in range(1, num_hops + 1): 47 | g.update_all(fn.u_mul_e(f'feat_{h - 1}', 'weight', 'msg'), fn.sum('msg', f'feat_{h}')) 48 | return [g.ndata.pop(f'feat_{h}') for h in range(num_hops + 1)] 49 | 50 | 51 | def train(args): 52 | set_random_seed(args.seed) 53 | g, labels, num_classes, train_idx, val_idx, test_idx = load_data(args.dataset, args.ogb_root) 54 | print('正在预先计算邻居聚集特征...') 55 | features = preprocess(g, g.ndata['feat'], args.num_hops) # List[tensor(N, d_in)]，长度为r+1 56 | train_feats = [feat[train_idx] for feat in features] 57 | val_feats = [feat[val_idx] for feat in features] 58 | test_feats = [feat[test_idx] for feat in features] 59 | 60 | model = SIGN( 61 | g.ndata['feat'].shape[1], args.num_hidden, num_classes, args.num_hops, 62 | args.num_layers, args.dropout 63 | ) 64 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 65 | 66 | for epoch in range(args.epochs): 67 | model.train() 68 | logits = model(train_feats) 69 | loss = F.cross_entropy(logits, labels[train_idx]) 70 | optimizer.zero_grad() 71 | loss.backward() 72 | optimizer.step() 73 | 74 | train_acc = accuracy(logits, labels[train_idx]) 75 | val_acc = evaluate(model, val_feats, labels[val_idx]) 76 | print('Epoch {:d} | Train Loss {:.4f} | Train Acc {:.4f} | Val Acc {:.4f}'.format( 77 | epoch, loss, train_acc, val_acc 78 | )) 79 | test_acc = evaluate(model, test_feats, labels[test_idx]) 80 | print('Test Acc {:.4f}'.format(test_acc)) 81 | 82 | 83 | def evaluate(model, feats, labels): 84 | model.eval() 85 | with torch.no_grad(): 86 | logits = model(feats) 87 | return accuracy(logits, labels) 88 | 89 | 90 | def main(): 91 | parser = argparse.ArgumentParser(description='SIGN') 92 | parser.add_argument('--seed', type=int, default=1, help='random seed') 93 | parser.add_argument( 94 | '--dataset', choices=['cora', 'reddit', 'amazon'], default='cora', help='dataset' 95 | ) 96 | parser.add_argument('--ogb-root', default='./ogb', help='root directory to OGB datasets') 97 | parser.add_argument('--num-hidden', type=int, default=256, help='number of hidden units') 98 | parser.add_argument('--num_hops', type=int, default=3, help='number of hops') 99 | parser.add_argument('--num-layers', type=int, default=2, help='number of feed-forward layers') 100 | parser.add_argument('--dropout', type=float, default=0.5, help='dropout probability') 101 | parser.add_argument('--epochs', type=int, default=50, help='number of training epochs') 102 | parser.add_argument('--lr', type=float, default=0.003, help='learning rate') 103 | parser.add_argument('--weight-decay', type=float, default=0., help='weight decay') 104 | args = parser.parse_args() 105 | print(args) 106 | train(args) 107 | 108 | 109 | if __name__ == '__main__': 110 | main() 111 | -------------------------------------------------------------------------------- /gnn/supergat/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/gnn/supergat/__init__.py -------------------------------------------------------------------------------- /gnn/supergat/attention.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import dgl.function as fn 4 | import torch 5 | import torch.nn as nn 6 | 7 | 8 | class GraphAttention(nn.Module): 9 | """图注意力模块，用于计算顶点邻居的重要性""" 10 | 11 | def __init__(self, out_dim, num_heads): 12 | super().__init__() 13 | self.out_dim = out_dim 14 | self.num_heads = num_heads 15 | 16 | def forward(self, g, feat_src, feat_dst): 17 | """ 18 | :param g: DGLGraph 同构图 19 | :param feat_src: tensor(N_src, K, d_out) 起点特征 20 | :param feat_dst: tensor(N_dst, K, d_out) 终点特征 21 | :return: tensor(E, K, 1) 所有顶点对的邻居重要性 22 | """ 23 | raise NotImplementedError 24 | 25 | 26 | class GATOriginalAttention(GraphAttention): 27 | """原始GAT注意力""" 28 | 29 | def __init__(self, out_dim, num_heads): 30 | super().__init__(out_dim, num_heads) 31 | self.attn_l = nn.Parameter(torch.FloatTensor(size=(1, num_heads, out_dim))) 32 | self.attn_r = nn.Parameter(torch.FloatTensor(size=(1, num_heads, out_dim))) 33 | self.reset_parameters() 34 | 35 | def reset_parameters(self): 36 | gain = nn.init.calculate_gain('relu') 37 | nn.init.xavier_normal_(self.attn_l, gain=gain) 38 | nn.init.xavier_normal_(self.attn_r, gain=gain) 39 | 40 | def forward(self, g, feat_src, feat_dst): 41 | el = (feat_src * self.attn_l).sum(dim=-1, keepdim=True) # (N_src, K, 1) 42 | er = (feat_dst * self.attn_r).sum(dim=-1, keepdim=True) # (N_dst, K, 1) 43 | g.srcdata['el'] = el 44 | g.dstdata['er'] = er 45 | g.apply_edges(fn.u_add_v('el', 'er', 'e')) 46 | return g.edata.pop('e') 47 | 48 | 49 | class DotProductAttention(GraphAttention): 50 | """点积注意力""" 51 | 52 | def forward(self, g, feat_src, feat_dst): 53 | g.srcdata['ft'] = feat_src 54 | g.dstdata['ft'] = feat_dst 55 | g.apply_edges(lambda edges: { 56 | 'e': torch.sum(edges.src['ft'] * edges.dst['ft'], dim=-1, keepdim=True) 57 | }) 58 | return g.edata.pop('e') 59 | 60 | 61 | class ScaledDotProductAttention(DotProductAttention): 62 | 63 | def forward(self, g, feat_src, feat_dst): 64 | return super().forward(g, feat_src, feat_dst) / math.sqrt(self.out_dim) 65 | 66 | 67 | class MixedGraphAttention(GATOriginalAttention, DotProductAttention): 68 | 69 | def forward(self, g, feat_src, feat_dst): 70 | return GATOriginalAttention.forward(self, g, feat_src, feat_dst) \ 71 | * torch.sigmoid(DotProductAttention.forward(self, g, feat_src, feat_dst)) 72 | 73 | 74 | GRAPH_ATTENTIONS = { 75 | 'GO': GATOriginalAttention, 76 | 'DP': DotProductAttention, 77 | 'SD': ScaledDotProductAttention, 78 | 'MX': MixedGraphAttention 79 | } 80 | 81 | 82 | def get_graph_attention(attn_type, out_dim, num_heads): 83 | if attn_type in GRAPH_ATTENTIONS: 84 | return GRAPH_ATTENTIONS[attn_type](out_dim, num_heads) 85 | else: 86 | raise ValueError('非法图注意力类型{}，可选项为{}'.format(attn_type, list(GRAPH_ATTENTIONS.keys()))) 87 | -------------------------------------------------------------------------------- /gnn/supergat/model.py: -------------------------------------------------------------------------------- 1 | """How to Find Your Friendly Neighborhood: Graph Attention Design with Self-Supervision (SuperGAT) 2 | 3 | 论文链接：https://openreview.net/pdf?id=Wi5KUNlqWty 4 | """ 5 | import dgl 6 | import dgl.function as fn 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | from dgl.ops import edge_softmax 11 | 12 | from gnn.supergat.attention import get_graph_attention 13 | from gnn.utils import RatioNegativeSampler 14 | 15 | 16 | class SuperGATConv(nn.Module): 17 | 18 | def __init__( 19 | self, in_dim, out_dim, num_heads, attn_type, neg_sample_ratio=0.5, 20 | feat_drop=0.0, attn_drop=0.0, negative_slope=0.2, activation=None): 21 | """SuperGAT层，自监督任务是连接预测 22 | 23 | :param in_dim: int 输入特征维数 24 | :param out_dim: int 输出特征维数 25 | :param num_heads: int 注意力头数K 26 | :param attn_type: str 注意力类型，可选择GO, DP, SD和MX 27 | :param neg_sample_ratio: float, optional 负样本边数量占正样本边数量的比例，默认0.5 28 | :param feat_drop: float, optional 输入特征Dropout概率，默认为0 29 | :param attn_drop: float, optional 注意力权重Dropout概率，默认为0 30 | :param negative_slope: float, optional LeakyReLU负斜率，默认为0.2 31 | :param activation: callable, optional 用于输出特征的激活函数，默认为None 32 | """ 33 | super().__init__() 34 | self.out_dim = out_dim 35 | self.num_heads = num_heads 36 | self.fc = nn.Linear(in_dim, num_heads * out_dim, bias=False) 37 | self.attn = get_graph_attention(attn_type, out_dim, num_heads) 38 | self.neg_sampler = RatioNegativeSampler(neg_sample_ratio) 39 | self.feat_drop = nn.Dropout(feat_drop) 40 | self.attn_drop = nn.Dropout(attn_drop) 41 | self.leaky_relu = nn.LeakyReLU(negative_slope) 42 | self.activation = activation 43 | 44 | def forward(self, g, feat): 45 | """ 46 | :param g: DGLGraph 同构图 47 | :param feat: tensor(N_src, d_in) 输入顶点特征 48 | :return: tensor(N_dst, K, d_out) 输出顶点特征 49 | """ 50 | with g.local_scope(): 51 | feat_src = self.fc(self.feat_drop(feat)).view(-1, self.num_heads, self.out_dim) 52 | feat_dst = feat_src[:g.num_dst_nodes()] if g.is_block else feat_src 53 | e = self.leaky_relu(self.attn(g, feat_src, feat_dst)) # (E, K, 1) 54 | g.edata['a'] = self.attn_drop(edge_softmax(g, e)) # (E, K, 1) 55 | g.srcdata['ft'] = feat_src 56 | # 消息传递 57 | g.update_all(fn.u_mul_e('ft', 'a', 'm'), fn.sum('m', 'ft')) 58 | out = g.dstdata['ft'] # (N_dst, K, d_out) 59 | 60 | if self.training: 61 | # 负采样 62 | neg_g = dgl.graph( 63 | self.neg_sampler(g, list(range(g.num_edges()))), num_nodes=g.num_nodes(), 64 | device=g.device 65 | ) 66 | neg_e = self.attn(neg_g, feat_src, feat_src) # (E', K, 1) 67 | self.attn_x = torch.cat([e, neg_e]).squeeze(dim=-1).mean(dim=1) # (E+E',) 68 | self.attn_y = torch.cat([torch.ones(e.shape[0]), torch.zeros(neg_e.shape[0])]) \ 69 | .to(self.attn_x.device) 70 | 71 | if self.activation: 72 | out = self.activation(out) 73 | return out 74 | 75 | def get_attn_loss(self): 76 | """返回自监督注意力损失（即连接预测损失）""" 77 | if self.training: 78 | return F.binary_cross_entropy_with_logits(self.attn_x, self.attn_y) 79 | else: 80 | return torch.tensor(0.) 81 | 82 | 83 | class SuperGAT(nn.Module): 84 | 85 | def __init__( 86 | self, in_dim, hidden_dim, out_dim, num_heads, attn_type, neg_sample_ratio=0.5, 87 | feat_drop=0.0, attn_drop=0.0, negative_slope=0.2): 88 | """两层SuperGAT模型 89 | 90 | :param in_dim: int 输入特征维数 91 | :param hidden_dim: int 隐含特征维数 92 | :param out_dim: int 输出特征维数 93 | :param num_heads: int 注意力头数K 94 | :param attn_type: str 注意力类型，可选择GO, DP, SD和MX 95 | :param neg_sample_ratio: float, optional 负样本边数量占正样本边数量的比例，默认0.5 96 | :param feat_drop: float, optional 输入特征Dropout概率，默认为0 97 | :param attn_drop: float, optional 注意力权重Dropout概率，默认为0 98 | :param negative_slope: float, optional LeakyReLU负斜率，默认为0.2 99 | """ 100 | super().__init__() 101 | self.conv1 = SuperGATConv( 102 | in_dim, hidden_dim, num_heads, attn_type, neg_sample_ratio, 103 | feat_drop, attn_drop, negative_slope, F.elu 104 | ) 105 | self.conv2 = SuperGATConv( 106 | num_heads * hidden_dim, out_dim, num_heads, attn_type, neg_sample_ratio, 107 | 0, attn_drop, negative_slope 108 | ) 109 | 110 | def forward(self, g, feat): 111 | """ 112 | :param g: DGLGraph 同构图 113 | :param feat: tensor(N, d_in) 输入顶点特征 114 | :return: tensor(N, d_out), tensor(1) 输出顶点特征和自监督注意力损失 115 | """ 116 | h = self.conv1(g, feat).flatten(start_dim=1) # (N, K, d_hid) -> (N, K*d_hid) 117 | h = self.conv2(g, h).mean(dim=1) # (N, K, d_out) -> (N, d_out) 118 | return h, self.conv1.get_attn_loss() + self.conv2.get_attn_loss() 119 | -------------------------------------------------------------------------------- /gnn/supergat/readme.txt: -------------------------------------------------------------------------------- 1 | Cora 2 | python -m gnn.supergat.train --dataset=cora 3 | Test Accuracy 0.8140 4 | 5 | CiteSeer 6 | python -m gnn.supergat.train --dataset=citeseer --epochs=80 7 | Test Accuracy 0.6930 8 | 9 | PubMed 10 | python -m gnn.supergat.train --dataset=pubmed 11 | Test Accuracy 0.7790 12 | 13 | ogbn-arxiv 14 | python -m gnn.supergat.train --dataset=ogbn-arxiv --ogb-root=/home/zzy/ogb/ --num-hidden=16 --dropout=0.2 --lr=0.05 15 | Test Accuracy 0.1101 16 | 17 | Cora-Full 18 | python -m gnn.supergat.train --dataset=cora_full --lr=0.01 19 | Test Accuracy 0.5175 20 | 21 | CS 22 | python -m gnn.supergat.train --dataset=cs --lr=0.01 23 | Test Accuracy 0.9162 24 | 25 | Physics 26 | python -m gnn.supergat.train --dataset=physics --lr=0.01 27 | Test Accuracy 0.9555 28 | 29 | Photo 30 | python -m gnn.supergat.train --dataset=photo --lr=0.01 31 | Test Accuracy 0.9255 32 | 33 | Computers 34 | python -m gnn.supergat.train --dataset=computers --lr=0.01 35 | Test Accuracy 0.8906 36 | -------------------------------------------------------------------------------- /gnn/supergat/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import torch 4 | import torch.nn.functional as F 5 | import torch.optim as optim 6 | from dgl.data import gnn_benckmark 7 | from ogb.nodeproppred import DglNodePropPredDataset 8 | 9 | from gnn.supergat.model import SuperGAT 10 | from gnn.utils import load_citation_dataset, split_idx, set_random_seed, accuracy 11 | 12 | DATASETS = [ 13 | 'cora', 'citeseer', 'pubmed', 'ogbn-arxiv', 14 | 'cora_full', 'cs', 'physics', 'photo', 'computers' 15 | ] 16 | 17 | 18 | def load_data(name, ogb_root, seed, device): 19 | if name == 'ogbn-arxiv': 20 | data = DglNodePropPredDataset('ogbn-arxiv', ogb_root) 21 | g, labels = data[0] 22 | split = data.get_idx_split() 23 | return g.to(device), labels.squeeze(dim=-1).to(device), data.num_classes, \ 24 | split['train'].to(device), split['valid'].to(device), split['test'].to(device) 25 | elif name in ('cora', 'citeseer', 'pubmed'): 26 | data = load_citation_dataset(name) 27 | elif name == 'cora_full': 28 | data = gnn_benckmark.CoraFullDataset() 29 | elif name in ('cs', 'physics'): 30 | data = gnn_benckmark.Coauthor(name) 31 | elif name in ('photo', 'computers'): 32 | data = gnn_benckmark.AmazonCoBuy(name) 33 | else: 34 | raise ValueError('Unknown dataset:', name) 35 | 36 | g = data[0].to(device) 37 | # https://github.com/dmlc/dgl/issues/2479 38 | num_classes = data.num_classes 39 | if name in ('photo', 'computers'): 40 | num_classes = g.ndata['label'].max().item() + 1 41 | if 'train_mask' in g.ndata: 42 | train_idx = g.ndata['train_mask'].nonzero(as_tuple=True)[0] 43 | val_idx = g.ndata['val_mask'].nonzero(as_tuple=True)[0] 44 | test_idx = g.ndata['test_mask'].nonzero(as_tuple=True)[0] 45 | else: 46 | train_idx, val_idx, test_idx = split_idx(torch.arange(g.num_nodes()), 0.2, 0.3, seed) 47 | return g, g.ndata['label'], num_classes, train_idx.to(device), val_idx.to(device), test_idx.to(device) 48 | 49 | 50 | def train(args): 51 | set_random_seed(args.seed) 52 | device = f'cuda:{args.device}' if torch.cuda.is_available() and args.device >= 0 else 'cpu' 53 | device = torch.device(device) 54 | 55 | g, labels, num_classes, train_idx, val_idx, test_idx = load_data(args.dataset, args.ogb_root, args.seed, device) 56 | features = g.ndata['feat'] 57 | 58 | model = SuperGAT( 59 | features.shape[1], args.num_hidden, num_classes, args.num_heads, args.attn_type, 60 | args.neg_sample_ratio, args.dropout, args.dropout 61 | ).to(device) 62 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 63 | for epoch in range(args.epochs): 64 | model.train() 65 | logits, attn_loss = model(g, features) 66 | loss = F.cross_entropy(logits[train_idx], labels[train_idx]) 67 | loss += args.attn_loss_weight * attn_loss 68 | optimizer.zero_grad() 69 | loss.backward() 70 | optimizer.step() 71 | 72 | train_acc = accuracy(logits[train_idx], labels[train_idx]) 73 | val_acc = evaluate(model, g, features, labels, val_idx) 74 | print('Epoch {:04d} | Loss {:.4f} | Train Acc {:.4f} | Val Acc {:.4f}'.format( 75 | epoch, loss.item(), train_acc, val_acc 76 | )) 77 | acc = evaluate(model, g, features, labels, test_idx) 78 | print('Test Accuracy {:.4f}'.format(acc)) 79 | 80 | 81 | def evaluate(model, g, features, labels, mask): 82 | model.eval() 83 | with torch.no_grad(): 84 | logits, _ = model(g, features) 85 | return accuracy(logits[mask], labels[mask]) 86 | 87 | 88 | def main(): 89 | parser = argparse.ArgumentParser(description='SuperGAT') 90 | parser.add_argument('--seed', type=int, default=1, help='random seed') 91 | parser.add_argument('--device', type=int, default=-1, help='GPU device') 92 | parser.add_argument('--dataset', choices=DATASETS, default='cora', help='dataset') 93 | parser.add_argument('--ogb-root', default='./ogb', help='root directory to OGB datasets') 94 | parser.add_argument('--num-hidden', type=int, default=8, help='number of hidden units') 95 | parser.add_argument('--num-heads', type=int, default=8, help='number of attention heads') 96 | parser.add_argument( 97 | '--attn-type', choices=['GO', 'DP', 'SD', 'MX'], default='MX', help='type of attention' 98 | ) 99 | parser.add_argument( 100 | '--neg-sample-ratio', type=float, default=0.8, 101 | help='ratio of the number of sampled negative edges to the number of positive edges' 102 | ) 103 | parser.add_argument('--dropout', type=float, default=0.6, help='attention dropout probability') 104 | parser.add_argument('--epochs', type=int, default=200, help='number of training epochs') 105 | parser.add_argument('--lr', type=float, default=0.005, help='learning rate') 106 | parser.add_argument('--weight-decay', type=float, default=5e-4, help='weight decay') 107 | parser.add_argument('--attn-loss-weight', type=float, default=2.0, help='attention loss weight') 108 | args = parser.parse_args() 109 | print(args) 110 | train(args) 111 | 112 | 113 | if __name__ == '__main__': 114 | main() 115 | -------------------------------------------------------------------------------- /gnn/utils/__init__.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | import dgl 4 | import numpy as np 5 | import torch 6 | 7 | from .data import * 8 | from .metapath import * 9 | from .metrics import * 10 | from .neg_sampler import * 11 | from .random_walk import * 12 | 13 | 14 | def set_random_seed(seed): 15 | """设置Python, numpy, PyTorch的随机数种子 16 | 17 | :param seed: int 随机数种子 18 | """ 19 | random.seed(seed) 20 | np.random.seed(seed) 21 | torch.manual_seed(seed) 22 | if torch.cuda.is_available(): 23 | torch.cuda.manual_seed(seed) 24 | dgl.seed(seed) 25 | 26 | 27 | def get_device(device): 28 | """返回指定的GPU设备 29 | 30 | :param device: int GPU编号，-1表示CPU 31 | :return: torch.device 32 | """ 33 | return torch.device(f'cuda:{device}' if device >= 0 and torch.cuda.is_available() else 'cpu') 34 | -------------------------------------------------------------------------------- /gnn/utils/data.py: -------------------------------------------------------------------------------- 1 | import dgl 2 | from dgl.data import citation_graph, rdf, knowledge_graph 3 | from dgl.utils import extract_node_subframes, set_new_frames 4 | from sklearn.model_selection import train_test_split 5 | 6 | __all__ = [ 7 | 'load_citation_dataset', 'load_rdf_dataset', 'load_kg_dataset', 8 | 'split_idx', 'add_reverse_edges' 9 | ] 10 | 11 | 12 | def load_citation_dataset(name): 13 | m = { 14 | 'cora': citation_graph.CoraGraphDataset, 15 | 'citeseer': citation_graph.CiteseerGraphDataset, 16 | 'pubmed': citation_graph.PubmedGraphDataset 17 | } 18 | try: 19 | return m[name]() 20 | except KeyError: 21 | raise ValueError('Unknown citation dataset: {}'.format(name)) 22 | 23 | 24 | def load_rdf_dataset(name): 25 | m = { 26 | 'aifb': rdf.AIFBDataset, 27 | 'mutag': rdf.MUTAGDataset, 28 | 'bgs': rdf.BGSDataset, 29 | 'am': rdf.AMDataset 30 | } 31 | try: 32 | return m[name]() 33 | except KeyError: 34 | raise ValueError('Unknown RDF dataset: {}'.format(name)) 35 | 36 | 37 | def load_kg_dataset(name): 38 | m = { 39 | 'wn18': knowledge_graph.WN18Dataset, 40 | 'FB15k': knowledge_graph.FB15kDataset, 41 | 'FB15k-237': knowledge_graph.FB15k237Dataset 42 | } 43 | try: 44 | return m[name]() 45 | except KeyError: 46 | raise ValueError('Unknown knowledge graph dataset: {}'.format(name)) 47 | 48 | 49 | def split_idx(samples, train_size, val_size, random_state=None): 50 | """将samples划分为训练集、测试集和验证集，需满足（用浮点数表示）： 51 | 52 | * 0 < train_size < 1 53 | * 0 < val_size < 1 54 | * train_size + val_size < 1 55 | 56 | :param samples: list/ndarray/tensor 样本集 57 | :param train_size: int or float 如果是整数则表示训练样本的绝对个数，否则表示训练样本占所有样本的比例 58 | :param val_size: int or float 如果是整数则表示验证样本的绝对个数，否则表示验证样本占所有样本的比例 59 | :param random_state: int, optional 随机数种子 60 | :return: (train, val, test) 类型与samples相同 61 | """ 62 | train, val = train_test_split(samples, train_size=train_size, random_state=random_state) 63 | if isinstance(val_size, float): 64 | val_size *= len(samples) / len(val) 65 | val, test = train_test_split(val, train_size=val_size, random_state=random_state) 66 | return train, val, test 67 | 68 | 69 | def add_reverse_edges(g): 70 | """给异构图的每种边添加反向边，返回新的异构图 71 | 72 | :param g: DGLGraph 异构图 73 | :return: DGLGraph 添加反向边之后的异构图 74 | """ 75 | data = {} 76 | for stype, etype, dtype in g.canonical_etypes: 77 | u, v = g.edges(etype=(stype, etype, dtype)) 78 | data[(stype, etype, dtype)] = u, v 79 | data[(dtype, etype + '_rev', stype)] = v, u 80 | new_g = dgl.heterograph(data, {ntype: g.num_nodes(ntype) for ntype in g.ntypes}) 81 | node_frames = extract_node_subframes(g, None) 82 | set_new_frames(new_g, node_frames=node_frames) 83 | return new_g 84 | -------------------------------------------------------------------------------- /gnn/utils/metapath.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | 3 | import dgl 4 | import torch 5 | 6 | __all__ = ['metapath_based_graph'] 7 | 8 | 9 | def metapath_based_graph(g, metapath, nids=None, edge_feat_name='inst'): 10 | """返回异构图基于给定元路径的邻居组成的图，元路径实例作为边特征，如果元路径是对称的则返回的是同构图，否则是二分图。 11 | 12 | 与dgl.metapath_reachable_graph()的区别：如果两个顶点之间有多个元路径实例则在返回的图中有多条边 13 | 14 | :param g: DGLGraph 异构图 15 | :param metapath: List[str or (str, str, str)] 元路径，边类型列表 16 | :param nids: tensor(N), optional 起点id，如果为None则选择该类型所有顶点 17 | :param edge_feat_name: str 保存元路径实例的边特征名称 18 | :return: DGLGraph 基于元路径的图 19 | """ 20 | instances = metapath_instances(g, metapath, nids) 21 | src_nodes, dst_nodes = instances[:, 0], instances[:, -1] 22 | src_type, dst_type = g.to_canonical_etype(metapath[0])[0], g.to_canonical_etype(metapath[-1])[2] 23 | if src_type == dst_type: 24 | mg = dgl.graph((src_nodes, dst_nodes), num_nodes=g.num_nodes(src_type)) 25 | mg.edata[edge_feat_name] = instances 26 | else: 27 | mg = dgl.heterograph( 28 | {(src_type, '_E', dst_type): (src_nodes, dst_nodes)}, 29 | {src_type: g.num_nodes(src_type), dst_type: g.num_nodes(dst_type)} 30 | ) 31 | mg.edges['_E'].data[edge_feat_name] = instances 32 | return mg 33 | 34 | 35 | def metapath_instances(g, metapath, nids=None): 36 | """返回异构图中给定元路径的所有实例。 37 | 38 | :param g: DGLGraph 异构图 39 | :param metapath: List[str or (str, str, str)] 元路径，边类型列表 40 | :param nids: tensor(N), optional 起点id，如果为None则选择该类型所有顶点 41 | :return: tensor(E, L) E为元路径实例个数，L为元路径长度 42 | """ 43 | src_type = g.to_canonical_etype(metapath[0])[0] 44 | if nids is None: 45 | nids = g.nodes(src_type) 46 | paths = nids.unsqueeze(1).tolist() 47 | for etype in metapath: 48 | new_paths = [] 49 | neighbors = etype_neighbors(g, etype) 50 | for path in paths: 51 | for neighbor in neighbors[path[-1]]: 52 | new_paths.append(path + [neighbor]) 53 | paths = new_paths 54 | return torch.tensor(paths, dtype=torch.long) 55 | 56 | 57 | def etype_neighbors(g, etype): 58 | """返回异构图中给定基于边类型的邻居。 59 | 60 | :param g: DGLGraph 异构图 61 | :param etype: (str, str, str) 规范边类型 62 | :return: Dict[int, List[int]] 每个源顶点基于该类型的边的邻居 63 | """ 64 | adj = g.adj(scipy_fmt='coo', etype=etype) 65 | neighbors = defaultdict(list) 66 | for u, v in zip(adj.row, adj.col): 67 | neighbors[u].append(v) 68 | return neighbors 69 | 70 | 71 | def to_ntype_list(g, metapath): 72 | """将边类型列表表示的元路径转换为顶点类型列表。 73 | 74 | 例如：['ap', 'pc', 'cp', 'pa] -> ['a', 'p', 'c', 'p', 'a'] 75 | 76 | :param g: DGLGraph 异构图 77 | :param metapath: List[str or (str, str, str)] 元路径，边类型列表 78 | :return: List[str] 元路径的顶点类型列表表示 79 | """ 80 | metapath = [g.to_canonical_etype(etype) for etype in metapath] 81 | return [metapath[0][0]] + [etype[2] for etype in metapath] 82 | 83 | 84 | def metapath_instance_feat(metapath, node_feats, instances): 85 | """返回元路径实例特征，由中间顶点的特征组成。 86 | 87 | :param metapath: List[str] 元路径，顶点类型列表 88 | :param node_feats: Dict[str, tendor(N_i, d)] 顶点类型到顶点特征的映射，所有类型顶点的特征应具有相同的维数d 89 | :param instances: tensor(E, L) 元路径实例，E为元路径实例个数，L为元路径长度 90 | :return: tensor(E, L, d) 元路径实例特征 91 | """ 92 | feat_dim = node_feats[metapath[0]].shape[1] 93 | inst_feat = torch.zeros(instances.shape + (feat_dim,)) 94 | for i, ntype in enumerate(metapath): 95 | inst_feat[:, i] = node_feats[ntype][instances[:, i]] 96 | return inst_feat 97 | 98 | 99 | def metapath_adj(g, metapath): 100 | """返回给定元路径连接的起点和终点类型的邻接矩阵。 101 | 102 | :param g: DGLGraph 103 | :param metapath: List[str or (str, str, str)] 元路径，边类型列表 104 | :return: scipy.sparse.csr_matrix 105 | """ 106 | adj = 1 107 | for etype in metapath: 108 | adj *= g.adj(etype=etype, scipy_fmt='csr') 109 | return adj 110 | -------------------------------------------------------------------------------- /gnn/utils/metrics.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from sklearn.cluster import KMeans 3 | from sklearn.metrics import f1_score, normalized_mutual_info_score, adjusted_rand_score 4 | 5 | __all__ = ['accuracy', 'micro_macro_f1_score', 'nmi_ari_score', 'mean_reciprocal_rank', 'hits_at'] 6 | 7 | 8 | def accuracy(logits, labels): 9 | """计算准确率 10 | 11 | :param logits: tensor(N, C) 预测概率，N为样本数，C为类别数 12 | :param labels: tensor(N) 正确标签 13 | :return: float 准确率 14 | """ 15 | return torch.sum(torch.argmax(logits, dim=1) == labels).item() * 1.0 / len(labels) 16 | 17 | 18 | def micro_macro_f1_score(logits, labels): 19 | """计算Micro-F1和Macro-F1得分 20 | 21 | :param logits: tensor(N, C) 预测概率，N为样本数，C为类别数 22 | :param labels: tensor(N) 正确标签 23 | :return: float, float Micro-F1和Macro-F1得分 24 | """ 25 | prediction = torch.argmax(logits, dim=1).long().numpy() 26 | labels = labels.numpy() 27 | micro_f1 = f1_score(labels, prediction, average='micro') 28 | macro_f1 = f1_score(labels, prediction, average='macro') 29 | return micro_f1, macro_f1 30 | 31 | 32 | def nmi_ari_score(logits, labels): 33 | """计算NMI和ARI得分 34 | 35 | :param logits: tensor(N, d) 预测概率，N为样本数 36 | :param labels: tensor(N) 正确标签 37 | :return: float, float NMI和ARI得分 38 | """ 39 | num_classes = logits.shape[1] 40 | prediction = KMeans(n_clusters=num_classes).fit_predict(logits.detach().numpy()) 41 | labels = labels.numpy() 42 | nmi = normalized_mutual_info_score(labels, prediction) 43 | ari = adjusted_rand_score(labels, prediction) 44 | return nmi, ari 45 | 46 | 47 | def mean_reciprocal_rank(predicts, answers): 48 | """计算平均倒数排名(MRR) = sum(1 / rank_i) / N，如果答案不在预测结果中，则排名按∞计算 49 | 50 | 例如，predicts=[[2, 0, 1], [2, 1, 0], [1, 0, 2]], answers=[1, 2, 8], 51 | ranks=[3, 2, ∞], MRR=(1/3+1/1+0)/3=4/9 52 | 53 | :param predicts: tensor(N, K) 预测结果，N为样本数，K为预测结果数 54 | :param answers: tensor(N) 正确答案的位置 55 | :return: float MRR∈[0, 1] 56 | """ 57 | ranks = torch.nonzero(predicts == answers.unsqueeze(1), as_tuple=True)[1] + 1 58 | return torch.sum(1.0 / ranks.float()).item() / len(predicts) 59 | 60 | 61 | def hits_at(n, predicts, answers): 62 | """计算Hits@n = #(rank_i <= n) / N 63 | 64 | 例如，predicts=[[2, 0, 1], [2, 1, 0], [1, 0, 2]], answers=[1, 2, 0], ranks=[3, 1, 2], Hits@2=2/3 65 | 66 | :param n: int 要计算答案排名前几的比例 67 | :param predicts: tensor(N, K) 预测结果，N为样本数，K为预测结果数 68 | :param answers: tensor(N) 正确答案的位置 69 | :return: float Hits@n∈[0, 1] 70 | """ 71 | ranks = torch.nonzero(predicts == answers.unsqueeze(1), as_tuple=True)[1] + 1 72 | return torch.sum(ranks <= n).float().item() / len(predicts) 73 | -------------------------------------------------------------------------------- /gnn/utils/neg_sampler.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | import torch 4 | from dgl.dataloading.negative_sampler import _BaseNegativeSampler 5 | 6 | __all__ = ['RatioNegativeSampler'] 7 | 8 | 9 | class RatioNegativeSampler(_BaseNegativeSampler): 10 | 11 | def __init__(self, neg_sample_ratio=1.0): 12 | """按一定正样本边的比例采样负样本边的负采样器 13 | 14 | :param neg_sample_ratio: float, optional 负样本边数量占正样本边数量的（大致）比例，默认为1 15 | """ 16 | self.neg_sample_ratio = neg_sample_ratio 17 | 18 | def _generate(self, g, eids, canonical_etype): 19 | stype, _, dtype = canonical_etype 20 | num_src_nodes, num_dst_nodes = g.num_nodes(stype), g.num_nodes(dtype) 21 | total = num_src_nodes * num_dst_nodes 22 | 23 | # 处理|Vs×Vd-E|= 0) 32 | -------------------------------------------------------------------------------- /kgrec/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/kgrec/__init__.py -------------------------------------------------------------------------------- /kgrec/kgcn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/kgrec/kgcn/__init__.py -------------------------------------------------------------------------------- /kgrec/kgcn/data.py: -------------------------------------------------------------------------------- 1 | import dgl 2 | import pandas as pd 3 | import torch 4 | from dgl.dataloading.negative_sampler import Uniform 5 | from sklearn.preprocessing import LabelEncoder 6 | 7 | 8 | class RatingKnowledgeGraphDataset: 9 | """基于知识图谱的用户评分数据集 10 | 11 | 读取用户-物品评分数据和知识图谱三元组，并分别构造为两个图： 12 | 13 | * user_item_graph: DGLGraph (user, rate, item)二分图，由正向（评分大于等于阈值）的交互关系组成 14 | * knowledge_graph: DGLGraph 同构图，其中0~N_item-1对应user_item_graph中的item顶点（即物品集合是实体集合的子集） 15 | ，边特征relation表示关系类型 16 | """ 17 | CONFIG = { 18 | 'movie': { 19 | 'rating_path': 'data/kgcn/movie/ratings.csv', 20 | 'rating_sep': ',', 21 | 'kg_path': 'data/kgcn/movie/kg.txt', 22 | 'item2id_path': 'data/kgcn/movie/item_index2entity_id.txt', 23 | 'threshold': 4.0 24 | }, 25 | 'music': { 26 | 'rating_path': 'data/kgcn/music/user_artists.dat', 27 | 'rating_sep': '\t', 28 | 'kg_path': 'data/kgcn/music/kg.txt', 29 | 'item2id_path': 'data/kgcn/music/item_index2entity_id.txt', 30 | 'threshold': 0.0 31 | } 32 | } 33 | 34 | def __init__(self, dataset): 35 | self.dataset = dataset 36 | cfg = self.CONFIG[dataset] 37 | 38 | rating = pd.read_csv( 39 | cfg['rating_path'], sep=cfg['rating_sep'], names=['user_id', 'item_id', 'rating'], 40 | usecols=[0, 1, 2], skiprows=1 41 | ) 42 | kg = pd.read_csv(cfg['kg_path'], sep='\t', names=['head', 'relation', 'tail']) 43 | item2entity = pd.read_csv(cfg['item2id_path'], sep='\t', names=['item_id', 'entity_id']) 44 | 45 | rating = rating[rating['item_id'].isin(item2entity['item_id'])] 46 | rating.reset_index(drop=True, inplace=True) 47 | rating['user_id'] = LabelEncoder().fit_transform(rating['user_id']) 48 | item2entity = dict(zip(item2entity['item_id'], item2entity['entity_id'])) 49 | rating['item_id'] = rating['item_id'].apply(item2entity.__getitem__) 50 | rating['label'] = rating['rating'].apply(lambda r: int(r >= cfg['threshold'])) 51 | rating = rating[rating['label'] == 1] 52 | user_item_graph = dgl.heterograph({ 53 | ('user', 'rate', 'item'): (rating['user_id'].to_numpy(), rating['item_id'].to_numpy()) 54 | }) 55 | 56 | # 负采样 57 | neg_sampler = Uniform(1) 58 | nu, nv = neg_sampler(user_item_graph, torch.arange(user_item_graph.num_edges())) 59 | u, v = user_item_graph.edges() 60 | self.user_item_graph = dgl.heterograph({('user', 'rate', 'item'): (torch.cat([u, nu]), torch.cat([v, nv]))}) 61 | self.user_item_graph.edata['label'] = torch.cat([torch.ones(u.shape[0]), torch.zeros(nu.shape[0])]) 62 | 63 | kg['relation'] = LabelEncoder().fit_transform(kg['relation']) 64 | # 有重边，即两个实体之间可能存在多条边，关系类型不同 65 | knowledge_graph = dgl.graph((kg['head'], kg['tail'])) 66 | knowledge_graph.edata['relation'] = torch.tensor(kg['relation'].tolist()) 67 | self.knowledge_graph = dgl.add_reverse_edges(knowledge_graph, copy_edata=True) 68 | 69 | def get_num(self): 70 | return self.user_item_graph.num_nodes('user'), self.knowledge_graph.num_nodes(), self.knowledge_graph.edata['relation'].max().item() + 1 71 | -------------------------------------------------------------------------------- /kgrec/kgcn/dataloader.py: -------------------------------------------------------------------------------- 1 | import dgl 2 | from dgl.dataloading import EdgeCollator, EdgeDataLoader 3 | from dgl.utils import prepare_tensor 4 | from torch.utils.data import DataLoader 5 | 6 | 7 | class KGCNEdgeCollator(EdgeCollator): 8 | 9 | def __init__(self, user_item_graph, eids, block_sampler, knowledge_graph, negative_sampler=None): 10 | """用于KGCN的EdgeCollator 11 | 12 | :param user_item_graph: DGLGraph 用户-物品图 13 | :param eids: tensor(E) 训练边id 14 | :param block_sampler: BlockSampler 邻居采样器 15 | :param knowledge_graph: DGLGraph 知识图谱 16 | :param negative_sampler: 负采样器 17 | """ 18 | super().__init__(user_item_graph, eids, block_sampler, knowledge_graph, negative_sampler=negative_sampler) 19 | 20 | def _collate(self, items): 21 | """根据边id采样子图 22 | 23 | :param items: tensor(B) 边id 24 | :return: tensor(N_src), DGLGraph, List[DGLBlock] 知识图谱的输入顶点id，用户-物品图的边子图， 25 | 知识图谱根据边id关联的物品id采样的多层MFG 26 | """ 27 | items = prepare_tensor(self.g_sampling, items, 'items') 28 | pair_graph = dgl.edge_subgraph(self.g, items) 29 | seed_nodes = pair_graph.ndata[dgl.NID] 30 | blocks = self.block_sampler.sample_blocks(self.g_sampling, seed_nodes['item']) 31 | input_nodes = blocks[0].srcdata[dgl.NID] 32 | return input_nodes, pair_graph, blocks 33 | 34 | def _collate_with_negative_sampling(self, items): 35 | """根据边id采样子图，并进行负采样 36 | 37 | :param items: tensor(B) 边id 38 | :return: tensor(N_src), DGLGraph, DGLGraph, List[DGLBlock] 知识图谱的输入顶点id，用户-物品图的边子图，负样本图， 39 | 知识图谱根据边id关联的物品id采样的多层MFG 40 | """ 41 | items = prepare_tensor(self.g_sampling, items, 'items') 42 | pair_graph = dgl.edge_subgraph(self.g, items, relabel_nodes=False) 43 | induced_edges = pair_graph.edata[dgl.EID] 44 | 45 | neg_srcdst = self.negative_sampler(self.g, items) 46 | neg_pair_graph = dgl.heterograph({self.g.canonical_etypes[0]: neg_srcdst}) 47 | 48 | pair_graph, neg_pair_graph = dgl.compact_graphs([pair_graph, neg_pair_graph]) 49 | pair_graph.edata[dgl.EID] = induced_edges 50 | seed_nodes = pair_graph.ndata[dgl.NID] 51 | 52 | blocks = self.block_sampler.sample_blocks(self.g_sampling, seed_nodes['item']) 53 | input_nodes = blocks[0].srcdata[dgl.NID] 54 | return input_nodes, pair_graph, neg_pair_graph, blocks 55 | 56 | 57 | class KGCNEdgeDataLoader(EdgeDataLoader): 58 | 59 | def __init__(self, user_item_graph, eids, block_sampler, knowledge_graph, negative_sampler=None, device='cpu', **kwargs): 60 | """用于KGCN的EdgeDataLoader 61 | 62 | :param user_item_graph: DGLGraph 用户-物品图 63 | :param eids: tensor(E) 训练边id 64 | :param block_sampler: BlockSampler 邻居采样器 65 | :param knowledge_graph: DGLGraph 知识图谱 66 | :param device: torch.device 67 | :param kwargs: DataLoader的其他参数 68 | """ 69 | super().__init__(user_item_graph, eids, block_sampler, device, **kwargs) 70 | self.collator = KGCNEdgeCollator(user_item_graph, eids, block_sampler, knowledge_graph, negative_sampler) 71 | self.dataloader = DataLoader(self.collator.dataset, collate_fn=self.collator.collate, **kwargs) 72 | -------------------------------------------------------------------------------- /kgrec/kgcn/model.py: -------------------------------------------------------------------------------- 1 | """Knowledge Graph Convolutional Networks for Recommender Systems (KGCN) 2 | 3 | 论文链接：https://arxiv.org/pdf/1904.12575 4 | """ 5 | import dgl 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | 11 | class KGCNLayer(nn.Module): 12 | 13 | def __init__(self, hidden_dim, aggregator): 14 | """KGCN层 15 | 16 | :param hidden_dim: int 隐含特征维数d 17 | :param aggregator: str 实体表示与邻居表示的组合方式：sum, concat, neighbor 18 | """ 19 | super().__init__() 20 | self.hidden_dim = hidden_dim 21 | self.aggregator = aggregator 22 | if aggregator == 'concat': 23 | self.w = nn.Linear(2 * hidden_dim, hidden_dim) 24 | else: 25 | self.w = nn.Linear(hidden_dim, hidden_dim) 26 | 27 | def forward(self, src_feat, dst_feat, rel_feat, user_feat, activation): 28 | """ 29 | :param src_feat: tensor(B, K^h, K, d) 输入实体表示，B为batch大小，K为邻居个数，h为跳步数/层数 30 | :param dst_feat: tensor(B, K^h, d) 目标实体表示 31 | :param rel_feat: tensor(B, K^h, K, d) 关系表示 32 | :param user_feat: tensor(B, d) 用户表示 33 | :param activation: callable 激活函数 34 | :return: tensor(B, K^h, d) 输出实体表示 35 | """ 36 | batch_size = user_feat.shape[0] 37 | user_feat = user_feat.view(batch_size, 1, 1, self.hidden_dim) # (B, 1, 1, d) 38 | 39 | user_rel_scores = (user_feat * rel_feat).sum(dim=-1) # (B, K^h, K) 40 | user_rel_scores = F.softmax(user_rel_scores, dim=-1).unsqueeze(dim=-1) # (B, K^h, K, 1) 41 | agg = (user_rel_scores * src_feat).sum(dim=2) # (B, K^h, d) 42 | 43 | if self.aggregator == 'sum': 44 | out = (dst_feat + agg).view(-1, self.hidden_dim) # (B*K^h, d) 45 | elif self.aggregator == 'concat': 46 | out = torch.cat([dst_feat, agg], dim=-1).view(-1, 2 * self.hidden_dim) # (B*K^h, 2d) 47 | else: 48 | out = agg.view(-1, self.hidden_dim) # (B*K^h, d) 49 | 50 | out = self.w(out).view(batch_size, -1, self.hidden_dim) # (B, K^h, d) 51 | return activation(out) 52 | 53 | 54 | class KGCN(nn.Module): 55 | 56 | def __init__(self, hidden_dim, neighbor_size, aggregator, num_hops, num_users, num_entities, num_rels): 57 | super().__init__() 58 | self.hidden_dim = hidden_dim 59 | self.neighbor_size = neighbor_size 60 | self.num_hops = num_hops 61 | self.aggregator = KGCNLayer(hidden_dim, aggregator) 62 | self.user_embed = nn.Embedding(num_users, hidden_dim) 63 | self.entity_embed = nn.Embedding(num_entities, hidden_dim) 64 | self.rel_embed = nn.Embedding(num_rels, hidden_dim) 65 | 66 | def forward(self, pair_graph, blocks): 67 | """ 68 | :param pair_graph: DGLGraph 用户-物品子图 69 | :param blocks: List[DGLBlock] 知识图谱的MFG，blocks[-1].dstnodes()对应items 70 | :return: tensor(B) 用户-物品预测概率 71 | """ 72 | u, v = pair_graph.edges() 73 | users = pair_graph.nodes['user'].data[dgl.NID][u] 74 | user_feat = self.user_embed(users) # (B, d) 75 | entities, relations = self._get_neighbors(v, blocks) 76 | item_feat = self._aggregate(entities, relations, user_feat) # (B, d) 77 | scores = (user_feat * item_feat).sum(dim=-1) 78 | return torch.sigmoid(scores) 79 | 80 | def _get_neighbors(self, v, blocks): 81 | batch_size = v.shape[0] 82 | entities, relations = [blocks[-1].dstdata[dgl.NID][v].unsqueeze(dim=-1)], [] 83 | for b in reversed(blocks): 84 | u, dst = b.in_edges(v) 85 | entities.append(b.srcdata[dgl.NID][u].view(batch_size, -1)) 86 | relations.append(b.edata['relation'][b.edge_ids(u, dst)].view(batch_size, -1)) 87 | v = u 88 | return entities, relations 89 | 90 | def _aggregate(self, entities, relations, user_feat): 91 | batch_size = user_feat.shape[0] 92 | entity_feats = [self.entity_embed(entity) for entity in entities] 93 | rel_feats = [self.rel_embed(rel) for rel in relations] 94 | for h in range(self.num_hops): 95 | activation = torch.tanh if h == self.num_hops - 1 else torch.sigmoid 96 | new_entity_feats = [ 97 | self.aggregator( 98 | entity_feats[i + 1].view(batch_size, -1, self.neighbor_size, self.hidden_dim), 99 | entity_feats[i], 100 | rel_feats[i].view(batch_size, -1, self.neighbor_size, self.hidden_dim), 101 | user_feat, activation 102 | ) for i in range(self.num_hops - h) 103 | ] 104 | entity_feats = new_entity_feats 105 | return entity_feats[0].view(batch_size, self.hidden_dim) 106 | -------------------------------------------------------------------------------- /kgrec/kgcn/readme.txt: -------------------------------------------------------------------------------- 1 | MovieLens-20M 2 | python -m kgrec.kgcn.train --dataset=movie --num-hidden=32 --aggregator=sum --num-hops=2 --epochs=5 --batch-size=65536 --neighbor-size=4 --lr=0.02 --weight-decay=0.0000001 3 | Train AUC 0.9742 | Train F1 0.9293 | Test AUC 0.9656 | Test F1 0.9175 4 | 作者代码结果： 5 | auc_score: 0.9903, f1_score: 0.9678 6 | 7 | Last.FM 8 | python -m kgrec.kgcn.train --dataset=music --num-hidden=16 --aggregator=sum --num-hops=1 --epochs=30 --batch-size=32 --neighbor-size=8 --lr=0.0005 --weight-decay=0.0001 9 | Train AUC 0.9793 | Train F1 0.9245 | Test AUC 0.8057 | Test F1 0.7243 10 | 作者代码结果： 11 | auc_score: 0.8590, f1_score: 0.7713 12 | 13 | （虽然论文中的公式看起来与普通的消息传递一样，但使用图上的消息传递并不容易实现，因此作者代码使用普通的张量计算实现，这也是必须为每个实体采样固定数量的邻居的原因） 14 | -------------------------------------------------------------------------------- /kgrec/kgcn/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import dgl 4 | import torch 5 | import torch.nn.functional as F 6 | import torch.optim as optim 7 | from dgl.dataloading import MultiLayerNeighborSampler 8 | from sklearn.metrics import roc_auc_score, f1_score 9 | from sklearn.model_selection import train_test_split 10 | 11 | from gnn.utils import set_random_seed, get_device 12 | from kgrec.kgcn.data import RatingKnowledgeGraphDataset 13 | from kgrec.kgcn.dataloader import KGCNEdgeDataLoader 14 | from kgrec.kgcn.model import KGCN 15 | 16 | 17 | def train(args): 18 | set_random_seed(args.seed) 19 | device = get_device(args.device) 20 | data = RatingKnowledgeGraphDataset(args.dataset) 21 | user_item_graph = data.user_item_graph 22 | knowledge_graph = dgl.sampling.sample_neighbors( 23 | data.knowledge_graph, data.knowledge_graph.nodes(), args.neighbor_size, replace=True 24 | ) 25 | 26 | train_eids, test_eids = train_test_split( 27 | torch.arange(user_item_graph.num_edges()), train_size=args.train_size, 28 | random_state=args.seed 29 | ) 30 | sampler = MultiLayerNeighborSampler([args.neighbor_size] * args.num_hops) 31 | train_loader = KGCNEdgeDataLoader( 32 | user_item_graph, train_eids, sampler, knowledge_graph, 33 | device=device, batch_size=args.batch_size 34 | ) 35 | test_loader = KGCNEdgeDataLoader( 36 | user_item_graph, test_eids, sampler, knowledge_graph, 37 | device=device, batch_size=args.batch_size 38 | ) 39 | 40 | model = KGCN(args.num_hidden, args.neighbor_size, args.aggregator, args.num_hops, *data.get_num()).to(device) 41 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) 42 | for epoch in range(args.epochs): 43 | model.train() 44 | losses = [] 45 | for _, pair_graph, blocks in train_loader: 46 | scores = model(pair_graph, blocks) 47 | loss = F.binary_cross_entropy(scores, pair_graph.edata['label']) 48 | losses.append(loss.item()) 49 | optimizer.zero_grad() 50 | loss.backward() 51 | optimizer.step() 52 | print('Epoch {:d} | Train Loss {:.4f} | Train AUC {:.4f} | Train F1 {:.4f} | Test AUC {:.4f} | Test F1 {:.4f}'.format( 53 | epoch, sum(losses) / len(losses), *evaluate(model, train_loader), *evaluate(model, test_loader) 54 | )) 55 | 56 | 57 | @torch.no_grad() 58 | def evaluate(model, loader): 59 | model.eval() 60 | auc_scores, f1_scores = [], [] 61 | for _, pair_graph, blocks in loader: 62 | u, v = pair_graph.edges() 63 | scores = model(pair_graph, blocks).detach().cpu().numpy() 64 | labels = pair_graph.edata['label'].int().cpu().numpy() 65 | auc_scores.append(roc_auc_score(labels, scores)) 66 | f1_scores.append(f1_score(labels, scores > 0.5)) 67 | return sum(auc_scores) / len(auc_scores), sum(f1_scores) / len(f1_scores) 68 | 69 | 70 | def main(): 71 | parser = argparse.ArgumentParser(description='KGCN Link Prediction') 72 | parser.add_argument('--seed', type=int, default=0, help='random seed') 73 | parser.add_argument('--device', type=int, default=0, help='GPU device') 74 | parser.add_argument('--dataset', choices=['movie', 'music'], default='music', help='dataset') 75 | parser.add_argument('--train-size', type=float, default=0.8, help='size of training dataset') 76 | parser.add_argument('--num-hidden', type=int, default=16, help='number of hidden units') 77 | parser.add_argument('--aggregator', choices=['sum', 'concat', 'neighbor'], help='aggregator') 78 | parser.add_argument('--num-hops', type=int, default=1, help='number of hops') 79 | parser.add_argument('--epochs', type=int, default=20, help='number of training epochs') 80 | parser.add_argument('--batch-size', type=int, default=32, help='batch size') 81 | parser.add_argument('--neighbor-size', type=int, default=8, help='number of sampled neighbors') 82 | parser.add_argument('--lr', type=float, default=5e-4, help='learning rate') 83 | parser.add_argument('--weight-decay', type=float, default=1e-4, help='weight decay') 84 | args = parser.parse_args() 85 | print(args) 86 | train(args) 87 | 88 | 89 | if __name__ == '__main__': 90 | main() 91 | -------------------------------------------------------------------------------- /nlp/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/nlp/__init__.py -------------------------------------------------------------------------------- /nlp/tfms/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZZy979/pytorch-tutorial/61d33441b15521bc7f0fc83244594c672af8de9e/nlp/tfms/__init__.py -------------------------------------------------------------------------------- /nlp/tfms/causal_lm_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from transformers import AutoModelForCausalLM, AutoTokenizer, top_k_top_p_filtering 4 | 5 | tokenizer = AutoTokenizer.from_pretrained('gpt2') 6 | model = AutoModelForCausalLM.from_pretrained('gpt2') 7 | 8 | sequence = 'Hugging Face is based in DUMBO, New York City, and ' 9 | input_ids = tokenizer.encode(sequence, return_tensors='pt') 10 | 11 | # get logits of last hidden state 12 | next_token_logits = model(input_ids).logits[:, -1, :] 13 | 14 | # filter 15 | filtered_next_token_logits = top_k_top_p_filtering(next_token_logits, top_k=50, top_p=1.0) 16 | 17 | # sample 18 | probs = F.softmax(filtered_next_token_logits, dim=-1) 19 | next_tokens = torch.multinomial(probs, num_samples=10)[0] 20 | for token in next_tokens: 21 | print(sequence + tokenizer.decode(token)) 22 | -------------------------------------------------------------------------------- /nlp/tfms/eqa_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from transformers import AutoTokenizer, AutoModelForQuestionAnswering 3 | 4 | name = 'bert-large-uncased-whole-word-masking-finetuned-squad' 5 | tokenizer = AutoTokenizer.from_pretrained(name) 6 | model = AutoModelForQuestionAnswering.from_pretrained(name) 7 | 8 | text = r'''🤗 Transformers (formerly known as pytorch-transformers and pytorch-pretrained-bert) provides general-purpose 9 | architectures (BERT, GPT-2, RoBERTa, XLM, DistilBert, XLNet…) for Natural Language Understanding (NLU) and Natural 10 | Language Generation (NLG) with over 32+ pretrained models in 100+ languages and deep interoperability between 11 | TensorFlow 2.0 and PyTorch.''' 12 | questions = [ 13 | 'How many pretrained models are available in 🤗 Transformers?', 14 | 'What does 🤗 Transformers provide?', 15 | '🤗 Transformers provides interoperability between which frameworks?', 16 | ] 17 | for question in questions: 18 | inputs = tokenizer(question, text, add_special_tokens=True, return_tensors='pt') 19 | input_ids = inputs['input_ids'].tolist()[0] 20 | 21 | text_tokens = tokenizer.convert_ids_to_tokens(input_ids) 22 | output = model(**inputs) 23 | answer_start_scores, answer_end_scores = output.start_logits, output.end_logits 24 | 25 | # Get the most likely beginning of answer with the argmax of the score 26 | answer_start = torch.argmax(answer_start_scores) 27 | # Get the most likely end of answer with the argmax of the score 28 | answer_end = torch.argmax(answer_end_scores) + 1 29 | answer = tokenizer.convert_tokens_to_string( 30 | tokenizer.convert_ids_to_tokens(input_ids[answer_start:answer_end]) 31 | ) 32 | print('Question:', question) 33 | print('Answer:', answer) 34 | -------------------------------------------------------------------------------- /nlp/tfms/eqa_pipeline.py: -------------------------------------------------------------------------------- 1 | from transformers import pipeline 2 | 3 | qa = pipeline('question-answering') 4 | context = '''Extractive Question Answering is the task of extracting an answer from a text given a question. 5 | An example of a question answering dataset is the SQuAD dataset, which is entirely based on that task. 6 | If you would like to fine-tune a model on a SQuAD task, you may leverage the examples/question-answering/run_squad.py script.''' 7 | questions = [ 8 | 'What is extractive question answering?', 9 | 'What is a good example of a question answering dataset?' 10 | ] 11 | for question in questions: 12 | result = qa(question=question, context=context) 13 | print('Question:', question) 14 | print('Answer: {0[answer]}, score: {0[score]}, start: {0[start]}, end: {0[end]}'.format(result)) 15 | -------------------------------------------------------------------------------- /nlp/tfms/masked_lm_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from transformers import AutoModelForMaskedLM, AutoTokenizer 3 | 4 | tokenizer = AutoTokenizer.from_pretrained('distilbert-base-cased') 5 | model = AutoModelForMaskedLM.from_pretrained('distilbert-base-cased') 6 | 7 | sequence = f'Distilled models are smaller than the models they mimic. Using them instead of the' \ 8 | f' large versions would help {tokenizer.mask_token} our carbon footprint.' 9 | print(sequence) 10 | 11 | inputs = tokenizer.encode(sequence, return_tensors='pt') 12 | mask_token_index = torch.where(inputs == tokenizer.mask_token_id)[1] 13 | 14 | token_logits = model(inputs).logits 15 | mask_token_logits = token_logits[0, mask_token_index, :] 16 | top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist() 17 | for token in top_5_tokens: 18 | print(sequence.replace(tokenizer.mask_token, tokenizer.decode([token]))) 19 | -------------------------------------------------------------------------------- /nlp/tfms/masked_lm_pipeline.py: -------------------------------------------------------------------------------- 1 | from pprint import pprint 2 | 3 | from transformers import pipeline 4 | 5 | fill_mask = pipeline('fill-mask') 6 | print(fill_mask.tokenizer.mask_token) 7 | masked_seq = f'HuggingFace is creating a {fill_mask.tokenizer.mask_token}' \ 8 | f' that the community uses to solve NLP tasks.' 9 | print(masked_seq) 10 | pprint(fill_mask(masked_seq)) 11 | -------------------------------------------------------------------------------- /nlp/tfms/seq_clf_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from transformers import AutoTokenizer, AutoModelForSequenceClassification 3 | 4 | tokenizer = AutoTokenizer.from_pretrained('bert-base-cased-finetuned-mrpc') 5 | model = AutoModelForSequenceClassification.from_pretrained('bert-base-cased-finetuned-mrpc') 6 | sentence0 = 'The company HuggingFace is based in New York City' 7 | sentence1 = 'Apples are especially bad for your health' 8 | sentence2 = "HuggingFace's headquarters are situated in Manhattan" 9 | paraphrase = tokenizer(sentence0, sentence2, return_tensors='pt') 10 | not_paraphrase = tokenizer(sentence0, sentence1, return_tensors='pt') 11 | 12 | paraphrase_logits = model(**paraphrase).logits 13 | not_paraphrase_logits = model(**not_paraphrase).logits 14 | print('paraphrase_logits:', paraphrase_logits) 15 | print('not_paraphrase_logits:', not_paraphrase_logits) 16 | 17 | paraphrase_results = torch.softmax(paraphrase_logits, dim=1) 18 | not_paraphrase_results = torch.softmax(not_paraphrase_logits, dim=-1) 19 | 20 | # Should be paraphrase 21 | print('paraphrase_results:', paraphrase_results) 22 | # Should not be paraphrase 23 | print('not_paraphrase_results:', not_paraphrase_results) 24 | -------------------------------------------------------------------------------- /nlp/tfms/seq_clf_pipeline.py: -------------------------------------------------------------------------------- 1 | from transformers import pipeline 2 | 3 | classifier = pipeline('sentiment-analysis') 4 | print(classifier('I hate you')) 5 | print(classifier('I love you')) 6 | -------------------------------------------------------------------------------- /nlp/tfms/text_gen_model.py: -------------------------------------------------------------------------------- 1 | from transformers import AutoModelForCausalLM, AutoTokenizer 2 | 3 | model = AutoModelForCausalLM.from_pretrained('xlnet-base-cased') 4 | tokenizer = AutoTokenizer.from_pretrained('xlnet-base-cased') 5 | 6 | # Padding text helps XLNet with short prompts - proposed by Aman Rusia in https://github.com/rusiaaman/XLNet-gen#methodology 7 | PADDING_TEXT = '''In 1991, the remains of Russian Tsar Nicholas II and his family 8 | (except for Alexei and Maria) are discovered. 9 | The voice of Nicholas's young son, Tsarevich Alexei Nikolaevich, narrates the 10 | remainder of the story. 1883 Western Siberia, 11 | a young Grigori Rasputin is asked by his father and a group of men to perform magic. 12 | Rasputin has a vision and denounces one of the men as a horse thief. Although his 13 | father initially slaps him for making such an accusation, Rasputin watches as the 14 | man is chased outside and beaten. Twenty years later, Rasputin sees a vision of 15 | the Virgin Mary, prompting him to become a priest. Rasputin quickly becomes famous, 16 | with people, even a bishop, begging for his blessing. ''' 17 | prompt = 'Today the weather is really nice and I am planning on ' 18 | inputs = tokenizer.encode(PADDING_TEXT + prompt, add_special_tokens=False, return_tensors='pt') 19 | prompt_length = len(tokenizer.decode( 20 | inputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True 21 | )) 22 | outputs = model.generate(inputs, max_length=250, do_sample=True, top_p=0.95, top_k=60) 23 | generated = prompt + tokenizer.decode(outputs[0])[prompt_length:] 24 | print(generated) 25 | -------------------------------------------------------------------------------- /nlp/tfms/text_gen_pipeline.py: -------------------------------------------------------------------------------- 1 | from transformers import pipeline 2 | 3 | text_generator = pipeline('text-generation') 4 | result = text_generator('As far as I am concerned, I will', max_length=50, do_sample=False) 5 | print(result[0]['generated_text']) 6 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | dgl>=0.7.0 2 | gensim>=3.8.3 3 | matplotlib>=3.3.1 4 | networkx>=2.5 5 | nltk>=3.5 6 | numpy>=1.20.1 7 | ogb>=1.2.5 8 | pandas>=1.2.2 9 | rdflib>=5.0.0 10 | scikit-learn>=0.24.1 11 | scipy>=1.6.1 12 | torch>=1.7.1 13 | torchvision>=0.8.2 14 | tqdm>=4.57.0 15 | transformers>=4.2.2 16 | --------------------------------------------------------------------------------