#基于点的任务
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeatures
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import torch
import torch.nn.functional as F
from torch.nn import Linear
from torch_geometric.nn import GCNConv

dataset = Planetoid(root='/data/Planetoid', name='Cora', transform=NormalizeFeatures())

data = dataset[0]
# print(data)

#可视化部分
def visualize(h,color):
    #降维操作
    z = TSNE(n_components=2).fit_transform(h.detach().cpu().numpy())
    plt.figure(figsize=(10,10))
    plt.xticks([])
    plt.yticks([])
    plt.scatter(z[:,0],z[:,1],c=color,cmap='Set2',s=70)
    plt.show()

#全连接网络
# class MLP(torch.nn.Module):
#     def __init__(self,hidden_channels):
#         super().__init__()
#         self.lin1 = Linear(dataset.num_features,hidden_channels)
#         self.lin2 = Linear(hidden_channels,dataset.num_classes)
#
#     def forward(self,x):
#         x = self.lin1(x)
#         x = x.relu()
#         x = F.dropout(x,p = 0.5,training=self.training)
#         x = self.lin2(x)
#         return x

#GCN网络
class GCN(torch.nn.Module):
    def __init__(self,hidden_channels):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features,hidden_channels)
        self.conv2 = GCNConv(hidden_channels,dataset.num_classes)

    def forward(self,x,edge_index):
        x = self.conv1(x,edge_index)
        x = x.relu()
        x = F.dropout(x,p = 0.5,training=self.training)
        x = self.conv2(x,edge_index)
        return x


model = GCN(hidden_channels=16)
# model.eval()
# out = model(data.x,data.edge_index)
# visualize(out,data.y)

# print(model)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=0.01,weight_decay=5e-4)

#训练部分
def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x,data.edge_index)
    loss = criterion(out[data.train_mask],data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss

#测试部分
def test():
    model.eval()
    out = model(data.x,data.edge_index)
    pred = out.argmax(dim=1)
    test_correct = pred[data.test_mask] == data.y[data.test_mask]
    test_acc = int(test_correct.sum()) / int(data.test_mask.sum())
    return test_acc

for epoch in range(1,201):
    loss = train()
    print(
        'Epoch: {:03d}, Loss: {:.4f}, Test: {:.4f}'.format(epoch, loss, test())
    )

#准确率计算
test_acc = test()
print('Test Accuracy: {:.4f}'.format(test_acc))

model.eval()
out = model(data.x,data.edge_index)
visualize(out,data.y)

结果如下：