125 lines
3.9 KiB
Python
125 lines
3.9 KiB
Python
# import torch
|
|
# import torch.nn as nn
|
|
# from torch_geometric.nn import MessagePassing
|
|
# from torch.nn.utils import spectral_norm
|
|
|
|
# class GCNConvLayer(MessagePassing):
|
|
# def __init__(self, in_channels, out_channels):
|
|
# super().__init__(aggr="add")
|
|
# self.lin_l = spectral_norm(nn.Linear(in_channels, out_channels, bias=False))
|
|
# self.lin_r = spectral_norm(nn.Linear(in_channels, out_channels, bias=False))
|
|
|
|
# def forward(self, x, edge_index, edge_attr, deg):
|
|
# out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg)
|
|
# out = self.lin_l(out)
|
|
# return out
|
|
|
|
# def message(self, x_j, edge_attr):
|
|
# return x_j * edge_attr
|
|
|
|
# def aggregate(self, inputs, index, deg):
|
|
# """
|
|
# TODO: add docstring.
|
|
# """
|
|
# out = super().aggregate(inputs, index)
|
|
# deg = deg + 1e-7
|
|
# return out / deg.view(-1, 1)
|
|
|
|
|
|
# class CorrectionNet(nn.Module):
|
|
# def __init__(self, hidden_dim=8, n_layers=1):
|
|
# super().__init__()
|
|
# # self.enc = GCNConvLayer(1, hidden_dim)
|
|
# self.enc = nn.Sequential(
|
|
# spectral_norm(nn.Linear(1, hidden_dim//2)),
|
|
# nn.GELU(),
|
|
# spectral_norm(nn.Linear(hidden_dim//2, hidden_dim)),
|
|
# )
|
|
# self.layers = torch.nn.ModuleList([GCNConvLayer(hidden_dim, hidden_dim) for _ in range(n_layers)])
|
|
# self.relu = nn.GELU()
|
|
|
|
# self.dec = nn.Sequential(
|
|
# spectral_norm(nn.Linear(hidden_dim, hidden_dim//2)),
|
|
# nn.GELU(),
|
|
# spectral_norm(nn.Linear(hidden_dim//2, 1)),
|
|
# )
|
|
|
|
# def forward(self, x, edge_index, edge_attr, deg,):
|
|
# # h = self.enc(x, edge_index, edge_attr, deg)
|
|
# # h = self.relu(self.enc(x))
|
|
# h = self.enc(x)
|
|
# for layer in self.layers:
|
|
# h = layer(h, edge_index, edge_attr, deg)
|
|
# # h = self.norm(h)
|
|
# h = self.relu(h)
|
|
# # out = self.dec(h, edge_index, edge_attr, deg)
|
|
# out = self.dec(h)
|
|
# return out
|
|
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch_geometric.nn import MessagePassing
|
|
from torch.nn.utils import spectral_norm
|
|
|
|
|
|
class CorrectionNet(MessagePassing):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
|
|
def __init__(self, hidden_dim=16):
|
|
super().__init__(aggr="add")
|
|
self.in_net = nn.Sequential(
|
|
spectral_norm(nn.Linear(1, hidden_dim // 2)),
|
|
nn.GELU(),
|
|
spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
|
|
)
|
|
|
|
self.out_net = nn.Sequential(
|
|
spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)),
|
|
nn.GELU(),
|
|
spectral_norm(nn.Linear(hidden_dim // 2, 1)),
|
|
)
|
|
|
|
self.lin_msg = spectral_norm(
|
|
nn.Linear(hidden_dim, hidden_dim, bias=False)
|
|
)
|
|
self.lin_update = spectral_norm(
|
|
nn.Linear(hidden_dim, hidden_dim, bias=False)
|
|
)
|
|
self.alpha = nn.Parameter(torch.tensor(0.0))
|
|
self.beta = nn.Parameter(torch.tensor(0.0))
|
|
|
|
def forward(self, x, edge_index, edge_attr, deg):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
x = self.in_net(x)
|
|
out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg)
|
|
return self.out_net(out)
|
|
|
|
def message(self, x_j, edge_attr):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
alpha = torch.sigmoid(self.alpha)
|
|
msg = x_j * edge_attr
|
|
msg = (1 - alpha) * msg + alpha * self.lin_msg(msg)
|
|
return msg
|
|
|
|
def update(self, aggr_out, x):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
beta = torch.sigmoid(self.beta)
|
|
return aggr_out * (1 - beta) + self.lin_msg(x) * beta
|
|
|
|
def aggregate(self, inputs, index, deg):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
out = super().aggregate(inputs, index)
|
|
deg = deg + 1e-7
|
|
return out / deg.view(-1, 1)
|