random changes
This commit is contained in:
Filippo Olivo
@@ -1,5 +1,5 @@
|
||||
__all__ = ["GraphFiniteDifference", "GatingGNO"]
|
||||
|
||||
from .finite_difference import GraphFiniteDifference
|
||||
from .learnable_finite_difference import GraphFiniteDifference
|
||||
from .local_gno import GatingGNO
|
||||
from .point_net import PointNet
|
||||
|
||||
@@ -1,25 +0,0 @@
|
||||
from pina.model import GraphNeuralOperator
|
||||
import torch
|
||||
from torch_geometric.data import Data
|
||||
|
||||
|
||||
class GNO(torch.nn.Module):
|
||||
def __init__(
|
||||
self, x_ch_node, f_ch_node, hidden, layers, edge_ch=0, out_ch=1
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
lifting_operator = torch.nn.Linear(x_ch_node + f_ch_node, hidden)
|
||||
self.gno = GraphNeuralOperator(
|
||||
lifting_operator=lifting_operator,
|
||||
projection_operator=torch.nn.Linear(hidden, out_ch),
|
||||
edge_features=edge_ch,
|
||||
n_layers=layers,
|
||||
internal_n_layers=2,
|
||||
shared_weights=False,
|
||||
)
|
||||
|
||||
def forward(self, x, c, edge_index, edge_attr):
|
||||
x = torch.cat([x, c], dim=-1)
|
||||
x = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
|
||||
return self.gno(x)
|
||||
@@ -9,32 +9,27 @@ class FiniteDifferenceStep(MessagePassing):
|
||||
TODO: add docstring.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
aggr: str = "add",
|
||||
normalize: bool = True,
|
||||
root_weight: float = 1.0,
|
||||
):
|
||||
def __init__(self, aggr: str = "add", root_weight: float = 1.0):
|
||||
super().__init__(aggr=aggr)
|
||||
|
||||
self.normalize = normalize
|
||||
assert (
|
||||
aggr == "add"
|
||||
), "Per somme pesate, l'aggregazione deve essere 'add'."
|
||||
self.root_weight = float(root_weight)
|
||||
|
||||
def forward(self, x, edge_index, edge_weight, deg):
|
||||
def forward(self, x, edge_index, edge_attr, deg, weight=1.0):
|
||||
"""
|
||||
TODO: add docstring.
|
||||
"""
|
||||
out = self.propagate(edge_index, x=x, edge_weight=edge_weight, deg=deg)
|
||||
out = self.propagate(
|
||||
edge_index, x=x, edge_attr=edge_attr, deg=deg, weight=weight
|
||||
)
|
||||
return out
|
||||
|
||||
def message(self, x_j, edge_weight):
|
||||
def message(self, x_j, edge_attr):
|
||||
"""
|
||||
TODO: add docstring.
|
||||
"""
|
||||
return edge_weight.view(-1, 1) * x_j
|
||||
return edge_attr.view(-1, 1) * x_j
|
||||
|
||||
def aggregate(self, inputs, index, deg):
|
||||
"""
|
||||
@@ -44,11 +39,12 @@ class FiniteDifferenceStep(MessagePassing):
|
||||
deg = deg + 1e-7
|
||||
return out / deg.view(-1, 1)
|
||||
|
||||
def update(self, aggr_out, x):
|
||||
def update(self, aggr_out, x, weight):
|
||||
"""
|
||||
TODO: add docstring.
|
||||
"""
|
||||
return self.root_weight * aggr_out + (1 - self.root_weight) * x
|
||||
print(weight)
|
||||
return weight * aggr_out + (1 - weight) * x
|
||||
|
||||
|
||||
class GraphFiniteDifference(nn.Module):
|
||||
@@ -56,24 +52,22 @@ class GraphFiniteDifference(nn.Module):
|
||||
TODO: add docstring.
|
||||
"""
|
||||
|
||||
def __init__(self, max_iters: int = 1000, threshold: float = 1e-4):
|
||||
def __init__(self, max_iters: int = 5000, threshold: float = 1e-4):
|
||||
"""
|
||||
TODO: add docstring.
|
||||
"""
|
||||
super().__init__()
|
||||
self.max_iters = max_iters
|
||||
self.threshold = threshold
|
||||
self.fd_step = FiniteDifferenceStep(
|
||||
aggr="add", normalize=True, root_weight=1.0
|
||||
)
|
||||
self.fd_step = FiniteDifferenceStep(aggr="add", root_weight=1.0)
|
||||
|
||||
@staticmethod
|
||||
def _compute_deg(edge_index, edge_weight, num_nodes):
|
||||
def _compute_deg(edge_index, edge_attr, num_nodes):
|
||||
"""
|
||||
TODO: add docstring.
|
||||
"""
|
||||
deg = torch.zeros(num_nodes, device=edge_index.device)
|
||||
deg = deg.scatter_add(0, edge_index[1], edge_weight)
|
||||
deg = deg.scatter_add(0, edge_index[1], edge_attr)
|
||||
return deg + 1e-7
|
||||
|
||||
@staticmethod
|
||||
@@ -84,19 +78,29 @@ class GraphFiniteDifference(nn.Module):
|
||||
return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze()
|
||||
|
||||
def forward(
|
||||
self, x, edge_index, edge_weight, c, boundary_mask, boundary_values
|
||||
self,
|
||||
x,
|
||||
edge_index,
|
||||
edge_attr,
|
||||
c,
|
||||
boundary_mask,
|
||||
boundary_values,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
TODO: add docstring.
|
||||
"""
|
||||
edge_attr = 1 / edge_attr[:, -1]
|
||||
c_ij = self._compute_c_ij(c, edge_index)
|
||||
edge_weight = edge_weight * c_ij
|
||||
deg = self._compute_deg(edge_index, edge_weight, x.size(0))
|
||||
edge_attr = edge_attr * c_ij
|
||||
deg = self._compute_deg(edge_index, edge_attr, x.size(0))
|
||||
conv_thres = self.threshold * torch.norm(x)
|
||||
for _i in tqdm(range(self.max_iters)):
|
||||
out = self.fd_step(x, edge_index, edge_weight, deg)
|
||||
weight = 1.0
|
||||
for _i in range(self.max_iters):
|
||||
out = self.fd_step(x, edge_index, edge_attr, deg, weight=weight)
|
||||
weight = weight * 0.9999
|
||||
out[boundary_mask] = boundary_values.unsqueeze(-1)
|
||||
if torch.norm(out - x) < conv_thres:
|
||||
break
|
||||
x = out
|
||||
return out
|
||||
return out, _i + 1
|
||||
|
||||
@@ -108,14 +108,14 @@ class MLP(torch.nn.Module):
|
||||
tmp_layers.append(self._output_dim)
|
||||
|
||||
self._layers = []
|
||||
self._LayerNorm = []
|
||||
self._batchnorm = []
|
||||
for i in range(len(tmp_layers) - 1):
|
||||
|
||||
self._layers.append(
|
||||
self.spect_norm(nn.Linear(tmp_layers[i], tmp_layers[i + 1]))
|
||||
)
|
||||
|
||||
self._LayerNorm.append(nn.LazyLayerNorm())
|
||||
self._batchnorm.append(nn.LazyBatchNorm1d())
|
||||
|
||||
if isinstance(func, list):
|
||||
self._functions = func
|
||||
@@ -124,7 +124,7 @@ class MLP(torch.nn.Module):
|
||||
|
||||
unique_list = []
|
||||
for layer, func, bnorm in zip(
|
||||
self._layers[:-1], self._functions, self._LayerNorm
|
||||
self._layers[:-1], self._functions, self._batchnorm
|
||||
):
|
||||
|
||||
unique_list.append(layer)
|
||||
@@ -208,7 +208,7 @@ class TNet(nn.Module):
|
||||
)
|
||||
|
||||
self._function = function()
|
||||
self._bn1 = nn.LazyLayerNorm()
|
||||
self._bn1 = nn.LazyBatchNorm1d()
|
||||
|
||||
def forward(self, X):
|
||||
"""Forward pass for T-Net
|
||||
@@ -299,9 +299,9 @@ class PointNet(nn.Module):
|
||||
self._tnet_feature = TNet(input_dim=64)
|
||||
|
||||
self._function = function()
|
||||
self._bn1 = nn.LazyLayerNorm()
|
||||
self._bn2 = nn.LazyLayerNorm()
|
||||
self._bn3 = nn.LazyLayerNorm()
|
||||
self._bn1 = nn.LazyBatchNorm1d()
|
||||
self._bn2 = nn.LazyBatchNorm1d()
|
||||
self._bn3 = nn.LazyBatchNorm1d()
|
||||
|
||||
def concat(self, embedding, input_):
|
||||
"""Returns concatenation of global and local features for Point-Net
|
||||
@@ -370,3 +370,205 @@ class PointNet(nn.Module):
|
||||
X = self._mlp4(X)
|
||||
|
||||
return X
|
||||
|
||||
|
||||
class ConvTNet(nn.Module):
|
||||
"""T-Net base class. Implementation of T-Network with convolutional layers.
|
||||
|
||||
Reference: Ali Kashefi et al. https://arxiv.org/abs/2208.13434
|
||||
"""
|
||||
|
||||
def __init__(self, input_dim):
|
||||
"""T-Net block constructor
|
||||
|
||||
:param input_dim: input dimension of point cloud
|
||||
:type input_dim: int
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
function = nn.Tanh
|
||||
self._function = function()
|
||||
|
||||
self._block1 = nn.Sequential(
|
||||
nn.Conv1d(input_dim, 64, 1),
|
||||
nn.BatchNorm1d(64),
|
||||
self._function,
|
||||
nn.Conv1d(64, 128, 1),
|
||||
nn.BatchNorm1d(128),
|
||||
self._function,
|
||||
nn.Conv1d(128, 1024, 1),
|
||||
nn.BatchNorm1d(1024),
|
||||
self._function,
|
||||
)
|
||||
|
||||
self._block2 = MLP(
|
||||
input_dim=1024,
|
||||
output_dim=input_dim * input_dim,
|
||||
layers=[512, 256],
|
||||
func=function,
|
||||
batch_norm=True,
|
||||
)
|
||||
|
||||
def forward(self, X):
|
||||
"""Forward pass for T-Net
|
||||
|
||||
:param X: input tensor, shape [batch, $input_{dim}$, N]
|
||||
with batch the batch size, N number of points and $input_{dim}$
|
||||
the input dimension of the point cloud.
|
||||
:type X: torch.tensor
|
||||
:return: output affine matrix transformation, shape
|
||||
[batch, $input_{dim} \times input_{dim}$] with batch
|
||||
the batch size and $input_{dim}$ the input dimension
|
||||
of the point cloud.
|
||||
:rtype: torch.tensor
|
||||
"""
|
||||
|
||||
batch, input_dim = X.shape[0], X.shape[1]
|
||||
|
||||
# encoding using first MLP
|
||||
X = self._block1(X)
|
||||
|
||||
# applying symmetric function to aggregate information (using max as default)
|
||||
X, _ = torch.max(X, dim=-1)
|
||||
|
||||
# decoding using third MLP
|
||||
X = self._block2(X)
|
||||
|
||||
return X.reshape(batch, input_dim, input_dim)
|
||||
|
||||
|
||||
class ConvPointNet(nn.Module):
|
||||
"""Point-Net base class. Implementation of Point Network for segmentation.
|
||||
|
||||
Reference: Ali Kashefi et al. https://arxiv.org/abs/2208.13434
|
||||
"""
|
||||
|
||||
def __init__(self, input_dim, output_dim, tnet=False):
|
||||
"""Point-Net block constructor
|
||||
|
||||
:param input_dim: input dimension of point cloud
|
||||
:type input_dim: int
|
||||
:param output_dim: output dimension of point cloud
|
||||
:type output_dim: int
|
||||
:param tnet: apply T-Net transformation, defaults to False
|
||||
:type tnet: bool, optional
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self._function = nn.Tanh()
|
||||
self._use_tnet = tnet
|
||||
|
||||
self._block1 = nn.Sequential(
|
||||
nn.Conv1d(input_dim, 64, 1),
|
||||
nn.BatchNorm1d(),
|
||||
self._function,
|
||||
nn.Conv1d(64, 64, 1),
|
||||
nn.BatchNorm1d(64),
|
||||
self._function,
|
||||
)
|
||||
|
||||
self._block2 = nn.Sequential(
|
||||
nn.Conv1d(64, 64, 1),
|
||||
nn.BatchNorm1d(64),
|
||||
self._function,
|
||||
nn.Conv1d(64, 128, 1),
|
||||
nn.BatchNorm1d(128),
|
||||
self._function,
|
||||
nn.Conv1d(128, 1024, 1),
|
||||
nn.BatchNorm1d(1024),
|
||||
self._function,
|
||||
)
|
||||
|
||||
self._block3 = nn.Sequential(
|
||||
nn.Conv1d(1088, 512, 1),
|
||||
nn.BatchNorm1d(512),
|
||||
self._function,
|
||||
nn.Conv1d(512, 256, 1),
|
||||
nn.BatchNorm1d(256),
|
||||
self._function,
|
||||
nn.Conv1d(256, 128, 1),
|
||||
nn.BatchNorm1d(128),
|
||||
self._function,
|
||||
)
|
||||
|
||||
self._block4 = nn.Conv1d(128, output_dim, 1)
|
||||
|
||||
if self._use_tnet:
|
||||
self._tnet_transform = ConvTNet(input_dim=input_dim)
|
||||
self._tnet_feature = ConvTNet(input_dim=64)
|
||||
|
||||
def concat(self, embedding, input_):
|
||||
"""
|
||||
Returns concatenation of global and local features for Point-Net
|
||||
|
||||
:param embedding: global features of Point-Net, shape [batch, $input_{dim}$]
|
||||
with batch the batch size and $input_{dim}$ the input dimension
|
||||
of the point cloud.
|
||||
:type embedding: torch.tensor
|
||||
:param input_: local features of Point-Net, shape [batch, N, $input_{dim}$]
|
||||
with batch the batch size, N number of points and $input_{dim}$
|
||||
the input dimension of the point cloud.
|
||||
:type input_: torch.tensor
|
||||
:return: concatenation vector, shape [batch, N, $input_{dim}$]
|
||||
with batch the batch size, N number of points and $input_{dim}$
|
||||
:rtype: torch.tensor
|
||||
"""
|
||||
n_points = input_.shape[-1]
|
||||
embedding = embedding.unsqueeze(2).repeat(1, 1, n_points)
|
||||
return torch.cat([embedding, input_], dim=1)
|
||||
|
||||
def forward(self, X):
|
||||
"""Forward pass for Point-Net
|
||||
|
||||
:param X: input tensor, shape [batch, N, $input_{dim}$]
|
||||
with batch the batch size, N number of points and $input_{dim}$
|
||||
the input dimension of the point cloud.
|
||||
:type X: torch.tensor
|
||||
:return: segmentation vector, shape [batch, N, $output_{dim}$]
|
||||
with batch the batch size, N number of points and $output_{dim}$
|
||||
the output dimension of the point cloud.
|
||||
:rtype: torch.tensor
|
||||
"""
|
||||
|
||||
# permuting indeces
|
||||
X = X.permute(0, 2, 1)
|
||||
|
||||
# using transform tnet if needed
|
||||
if self._use_tnet:
|
||||
transform = self._tnet_transform(X)
|
||||
X = X.transpose(2, 1)
|
||||
X = torch.matmul(X, transform)
|
||||
X = X.transpose(2, 1)
|
||||
|
||||
# encoding using first MLP
|
||||
X = self._block1(X)
|
||||
|
||||
# using transform tnet if needed
|
||||
if self._use_tnet:
|
||||
transform = self._tnet_feature(X)
|
||||
X = X.transpose(2, 1)
|
||||
X = torch.matmul(X, transform)
|
||||
X = X.transpose(2, 1)
|
||||
|
||||
# saving latent representation for later concatanation
|
||||
latent = X
|
||||
|
||||
# encoding using second MLP
|
||||
X = self._block2(X)
|
||||
|
||||
# applying symmetric function to aggregate information (using max as default)
|
||||
X, _ = torch.max(X, dim=-1)
|
||||
|
||||
# concatenating with latent vector
|
||||
X = self.concat(X, latent)
|
||||
|
||||
# decoding using third MLP
|
||||
X = self._block3(X)
|
||||
|
||||
# decoding using fourth MLP
|
||||
X = self._block4(X)
|
||||
|
||||
# permuting indeces
|
||||
X = X.permute(0, 2, 1)
|
||||
|
||||
return X
|
||||
|
||||
Reference in New Issue
Block a user