transfer files
This commit is contained in:
Filippo Olivo
@@ -1,119 +1,22 @@
|
||||
# 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
|
||||
from torch_geometric.nn.conv import GCNConv
|
||||
|
||||
|
||||
class CorrectionNet(MessagePassing):
|
||||
"""
|
||||
TODO: add docstring.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_dim=16):
|
||||
class GCNConvLayer(MessagePassing):
|
||||
def __init__(self, in_channels, out_channels):
|
||||
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))
|
||||
self.lin_l = nn.Linear(in_channels, out_channels, bias=True)
|
||||
# self.lin_r = spectral_norm(nn.Linear(in_channels, out_channels, bias=False))
|
||||
|
||||
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)
|
||||
out = self.lin_l(out)
|
||||
return 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
|
||||
return x_j * edge_attr.view(-1, 1)
|
||||
|
||||
def aggregate(self, inputs, index, deg):
|
||||
"""
|
||||
@@ -122,3 +25,45 @@ class CorrectionNet(MessagePassing):
|
||||
out = super().aggregate(inputs, index)
|
||||
deg = deg + 1e-7
|
||||
return out / deg.view(-1, 1)
|
||||
|
||||
|
||||
class CorrectionNet(nn.Module):
|
||||
def __init__(self, input_dim=1, output_dim=1, hidden_dim=8, n_layers=8):
|
||||
super().__init__()
|
||||
self.enc = nn.Linear(input_dim, hidden_dim, bias=False)
|
||||
# self.layers = n_layers
|
||||
# self.l = GCNConv(hidden_dim, hidden_dim, aggr="mean")
|
||||
self.layers = torch.nn.ModuleList(
|
||||
[GCNConv(hidden_dim, hidden_dim, aggr="mean", bias=False) for _ in range(n_layers)]
|
||||
)
|
||||
self.dec = nn.Linear(hidden_dim, output_dim)
|
||||
|
||||
def forward(self, x, edge_index, edge_attr,):
|
||||
h = self.enc(x)
|
||||
# h = self.relu(h)
|
||||
for l in self.layers:
|
||||
# print(f"Forward pass layer {_}")
|
||||
h = l(h, edge_index, edge_attr)
|
||||
# h = self.relu(h)
|
||||
out = self.dec(h)
|
||||
return out
|
||||
|
||||
|
||||
class MLPNet(nn.Module):
|
||||
def __init__(self, input_dim=1, output_dim=1, hidden_dim=8, n_layers=1):
|
||||
super().__init__()
|
||||
layers = []
|
||||
func = torch.nn.ReLU
|
||||
|
||||
self.network = nn.Sequential(
|
||||
nn.Linear(input_dim, hidden_dim),
|
||||
func(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
func(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
func(),
|
||||
nn.Linear(hidden_dim, output_dim),
|
||||
)
|
||||
|
||||
def forward(self, x, edge_index=None, edge_attr=None):
|
||||
return self.network(x)
|
||||
Reference in New Issue
Block a user