rk-random/PACT-Net

1import torch2import torch.nn as nn3import torch.nn.functional as F4from torch_geometric.nn import GCNConv, GINConv, GATConv, SAGEConv, global_mean_pool5 6 7class GCN(torch.nn.Module):8    def __init__(self, in_channels, hidden_channels, out_channels):9        super(GCN, self).__init__()10        self.conv1 = GCNConv(in_channels, hidden_channels)11        self.conv2 = GCNConv(hidden_channels, hidden_channels)12        self.lin = nn.Linear(hidden_channels, out_channels)13 14    def forward(self, data):15        x, edge_index, batch = data.x, data.edge_index, data.batch16        x = F.relu(self.conv1(x, edge_index))17        x = F.relu(self.conv2(x, edge_index))18        x = global_mean_pool(x, batch)19        return self.lin(x)20 21 22class GIN(torch.nn.Module):23    def __init__(self, in_channels, hidden_channels, out_channels):24        super(GIN, self).__init__()25        nn1 = nn.Sequential(26            nn.Linear(in_channels, hidden_channels),27            nn.ReLU(),28            nn.Linear(hidden_channels, hidden_channels),29        )30        nn2 = nn.Sequential(31            nn.Linear(hidden_channels, hidden_channels),32            nn.ReLU(),33            nn.Linear(hidden_channels, hidden_channels),34        )35        self.conv1 = GINConv(nn1)36        self.conv2 = GINConv(nn2)37        self.lin = nn.Linear(hidden_channels, out_channels)38 39    def forward(self, data):40        x, edge_index, batch = data.x, data.edge_index, data.batch41        x = F.relu(self.conv1(x, edge_index))42        x = F.relu(self.conv2(x, edge_index))43        x = global_mean_pool(x, batch)44        return self.lin(x)45 46 47class GAT(torch.nn.Module):48    def __init__(self, in_channels, hidden_channels, out_channels, heads=4):49        super(GAT, self).__init__()50        self.conv1 = GATConv(in_channels, hidden_channels, heads=heads)51        self.conv2 = GATConv(hidden_channels * heads, hidden_channels, heads=1)52        self.lin = nn.Linear(hidden_channels, out_channels)53 54    def forward(self, data):55        x, edge_index, batch = data.x, data.edge_index, data.batch56        x = F.elu(self.conv1(x, edge_index))57        x = F.elu(self.conv2(x, edge_index))58        x = global_mean_pool(x, batch)59        return self.lin(x)60 61 62class GraphSAGE(torch.nn.Module):63    def __init__(self, in_channels, hidden_channels, out_channels):64        super(GraphSAGE, self).__init__()65        self.conv1 = SAGEConv(in_channels, hidden_channels)66        self.conv2 = SAGEConv(hidden_channels, hidden_channels)67        self.lin = nn.Linear(hidden_channels, out_channels)68 69    def forward(self, data):70        x, edge_index, batch = data.x, data.edge_index, data.batch71        x = F.relu(self.conv1(x, edge_index))72        x = F.relu(self.conv2(x, edge_index))73        x = global_mean_pool(x, batch)74        return self.lin(x)75