baseline_model.py

import os.path as osp
import argparse

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T
from torch_geometric.nn import GCNConv, ChebConv, GATConv


class GCN(torch.nn.Module):
    def __init__(self, dataset, hidden):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(dataset.num_features, hidden, cached=True)
        self.conv2 = GCNConv(hidden, int(dataset.num_classes), cached=True)
        # self.conv1 = ChebConv(data.num_features, 16, K=2)
        # self.conv2 = ChebConv(16, data.num_features, K=2)

    def forward(self, data):
        x, edge_index, edge_weight = data.x, data.edge_index, data.edge_attr
        x = F.relu(self.conv1(x, edge_index, edge_weight))
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index, edge_weight)
        return x

class ChebNet(torch.nn.Module):
    def __init__(self, dataset, hidden):
        super(ChebNet, self).__init__()
        # self.conv1 = GCNConv(dataset.num_features, hidden, cached=True)
        # self.conv2 = GCNConv(hidden, int(dataset.num_classes), cached=True)
        self.conv1 = ChebConv(dataset.num_features, hidden, K=2)
        self.conv2 = ChebConv(hidden, int(dataset.num_classes), K=2)

    def forward(self, data):
        x, edge_index, edge_weight = data.x, data.edge_index, data.edge_attr
        x = F.relu(self.conv1(x, edge_index, edge_weight))
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index, edge_weight)
        return x


class GAT(torch.nn.Module):
    def __init__(self, dataset, hidden):
        super(GAT, self).__init__()

        self.conv1 = GATConv(dataset.num_features, hidden, heads=8, dropout=0.6)
        # On the Pubmed dataset, use heads=8 in conv2.
        self.conv2 = GATConv(8 * hidden, int(dataset.num_classes), heads=1, concat=False,
                             dropout=0.6)

    def forward(self, data):
        x, edge_index, edge_weight = data.x, data.edge_index, data.edge_attr
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, data.edge_index)
        # return F.log_softmax(x, dim=-1)
        return x