请帮我写出基于PYG搭建的GIN实现图分类的代码

好的，以下是一个基于PYG搭建的GIN实现图分类的代码：

import torch
import torch.nn.functional as F
from torch.nn import Linear
from torch_geometric.nn import GINConv, global_add_pool

class GIN(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super(GIN, self).__init__()
        
        self.conv1 = GINConv(Linear(num_features, 32), eps=0)
        self.conv2 = GINConv(Linear(32, 64), eps=0)
        self.conv3 = GINConv(Linear(64, 128), eps=0)
        self.lin = Linear(128, num_classes)

    def forward(self, x, edge_index, batch):
        x = F.relu(self.conv1(x, edge_index))
        x = F.relu(self.conv2(x, edge_index))
        x = F.relu(self.conv3(x, edge_index))
        x = global_add_pool(x, batch)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin(x)
        return F.log_softmax(x, dim=-1)

这是一个三层的GIN模型，每一层都是使用线性层作为MLP，没有使用任何的激活函数来保证可微性。在最后一层使用了全局池化来将整个图信息压缩到一个向量中，然后再使用线性层进行分类。在模型训练过程中，我们使用了Dropout来防止过拟合，同时使用了Log Softmax作为输出层的激活函数。