rk-random/PACT-Net

1import torch2import torch_geometric3import torch.nn as nn4import torch.nn.functional as F5from torch_geometric.nn import (6    PNAConv,7    global_mean_pool,8    global_max_pool,9    global_add_pool,10)11from torch_geometric.utils import degree12 13 14class PolyatomicNet(nn.Module):15    def __init__(16        self,17        node_feat_dim,18        edge_feat_dim,19        graph_feat_dim,20        deg,21        hidden_dim=128,22        num_layers=5,23        dropout=0.1,24    ):25        super().__init__()26        self.graph_feat_dim = graph_feat_dim27        self.node_emb = nn.Linear(node_feat_dim, hidden_dim)28        self.deg = deg29        self.virtualnode_emb = nn.Embedding(1, hidden_dim)30        self.vn_mlp = nn.Sequential(31            nn.Linear(hidden_dim, hidden_dim),32            nn.ReLU(),33            nn.Linear(hidden_dim, hidden_dim),34        )35 36        # For graph-level feature projection37        self.graph_proj = nn.Sequential(38            nn.Linear(graph_feat_dim, hidden_dim),39            nn.ReLU(),40            nn.Linear(hidden_dim, hidden_dim),41        )42 43        # PNAConv requires degree preprocessing44        self.deg_emb = nn.Embedding(20, hidden_dim)  # cap degree buckets45 46        aggregators = ["mean", "min", "max", "std"]47        scalers = ["identity", "amplification", "attenuation"]48 49        self.convs = nn.ModuleList()50        self.bns = nn.ModuleList()51 52        for _ in range(num_layers):53            conv = PNAConv(54                in_channels=hidden_dim,55                out_channels=hidden_dim,56                aggregators=aggregators,57                scalers=scalers,58                edge_dim=edge_feat_dim,59                towers=4,60                pre_layers=1,61                post_layers=1,62                divide_input=True,63                deg=deg,64            )65            self.convs.append(conv)66            self.bns.append(nn.BatchNorm1d(hidden_dim))67 68        self.dropout = nn.Dropout(dropout)69 70        # Final readout71        self.readout = nn.Sequential(72            nn.Linear(hidden_dim * 3, hidden_dim),73            nn.ReLU(),74            nn.Dropout(dropout),75            nn.Linear(hidden_dim, hidden_dim // 2),76            nn.ReLU(),77            nn.Linear(hidden_dim // 2, 1),78        )79 80    def forward(self, data):81        x, edge_index, edge_attr, batch = (82            data.x,83            data.edge_index,84            data.edge_attr,85            data.batch,86        )87 88        deg = degree(edge_index[0], x.size(0), dtype=torch.long).clamp(max=19)89        h = self.node_emb(x) + self.deg_emb(deg)90 91        vn = self.virtualnode_emb(92            torch.zeros(batch.max().item() + 1, dtype=torch.long, device=x.device)93        )94 95        for conv, bn in zip(self.convs, self.bns):96            h = h + vn[batch]97            h = conv(h, edge_index, edge_attr)98            h = bn(h)99            h = F.relu(h)100            h = self.dropout(h)101            vn = vn + self.vn_mlp(global_mean_pool(h, batch))102 103        mean_pool = global_mean_pool(h, batch)104        max_pool = global_max_pool(h, batch)105        # add_pool = global_add_pool(h, batch)106 107        max_feat_dim = self.graph_feat_dim108 109        if hasattr(data, "graph_feats") and isinstance(110            data, torch_geometric.data.Batch  # type: ignore111        ):112            g_proj_list = []113            for g in data.to_data_list():114                g_feat = g.graph_feats.to(x.device)115 116                if g_feat.size(0) < max_feat_dim:117                    padded = torch.zeros(max_feat_dim, device=g_feat.device)118                    padded[: g_feat.size(0)] = g_feat119                    g_feat = padded120                elif g_feat.size(0) > max_feat_dim:121                    g_feat = g_feat[:max_feat_dim]122                g_feat = torch.nan_to_num(g_feat, nan=0.0, posinf=1e5, neginf=-1e5)123                g_proj_list.append(self.graph_proj(g_feat))124 125            g_proj = torch.stack(g_proj_list, dim=0)126 127        else:128            g_feat = data.graph_feats.to(x.device)129            if g_feat.size(0) < max_feat_dim:130                padded = torch.zeros(max_feat_dim, device=g_feat.device)131                padded[: g_feat.size(0)] = g_feat132                g_feat = padded133            elif g_feat.size(0) > max_feat_dim:134                g_feat = g_feat[:max_feat_dim]135            g_feat = torch.nan_to_num(g_feat, nan=0.0, posinf=1e5, neginf=-1e5)136            g_proj = self.graph_proj(g_feat).unsqueeze(0)137 138        final_input = torch.cat([mean_pool, max_pool, g_proj], dim=1)139        return self.readout(final_input).view(-1)140