folivo/thermal-conduction-ml
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Gradient accumulation in BPTT (#2)

Browse Source
This commit is contained in:
Filippo Olivo
2025-11-11 20:14:28 +01:00
committed by GitHub
parent 195c66b444
commit a2dd348423
4 changed files with 292 additions and 179 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
ThermalSolver/model/__init__.py
View File

@@ -1,13 +1,13 @@
__all__ = [
"GraphFiniteDifference",
# "GraphFiniteDifference",
"GatingGNO",
"LearnableGraphFiniteDifference",
# "LearnableGraphFiniteDifference",
"PointNet",
]
from .learnable_finite_difference import (
GraphFiniteDifference as LearnableGraphFiniteDifference,
)
from .finite_difference import GraphFiniteDifference as GraphFiniteDifference
# from .learnable_finite_difference import (
# GraphFiniteDifference as LearnableGraphFiniteDifference,
# )
# from .finite_difference import GraphFiniteDifference as GraphFiniteDifference
from .local_gno import GatingGNO
from .point_net import PointNet

6
ThermalSolver/model/finite_difference.py
View File

@@ -14,7 +14,7 @@ class FiniteDifferenceStep(MessagePassing):
aggr == "add"
), "Per somme pesate, l'aggregazione deve essere 'add'."
# self.root_weight = float(root_weight)
self.p = torch.nn.Parameter(torch.tensor(0.8))
self.p = torch.nn.Parameter(torch.tensor(1.0))
self.a = root_weight
def forward(self, x, edge_index, edge_attr, deg):
@@ -43,9 +43,7 @@ class FiniteDifferenceStep(MessagePassing):
"""
TODO: add docstring.
"""
a = torch.clamp(self.a, 0.0, 1.0)
return a * aggr_out + (1 - a) * x
# return self.a * aggr_out + (1 - self.a) * x
return aggr_out
class GraphFiniteDifference(nn.Module):

179
ThermalSolver/model/learnable_finite_difference.py
View File

@@ -2,6 +2,40 @@ import torch
import torch.nn as nn
from torch_geometric.nn import MessagePassing
from torch.nn.utils import spectral_norm
from matplotlib.tri import Triangulation
from matplotlib import pyplot as plt
def _plot_mesh(y_pred, batch, iteration=None):
idx = batch.batch == 0
y = batch.y[idx].detach().cpu()
y_pred = y_pred[idx].detach().cpu()
pos = batch.pos[idx].detach().cpu()
pos = pos.detach().cpu()
tria = Triangulation(pos[:, 0], pos[:, 1])
plt.figure(figsize=(18, 5))
plt.subplot(1, 3, 1)
plt.tricontourf(tria, y.squeeze().numpy(), levels=14)
plt.colorbar()
plt.title("True temperature")
plt.subplot(1, 3, 2)
plt.tricontourf(tria, y_pred.squeeze().numpy(), levels=14)
plt.colorbar()
plt.title("Predicted temperature")
plt.subplot(1, 3, 3)
plt.tricontourf(tria, torch.abs(y_pred - y).squeeze().numpy(), levels=14)
plt.colorbar()
plt.title("Error")
plt.suptitle("GNO", fontsize=16)
name = (
f"images/gno_iter_{iteration:04d}.png"
if iteration is not None
else "gno.png"
)
plt.savefig(name, dpi=72)
plt.close()
class FiniteDifferenceStep(MessagePassing):
@@ -9,50 +43,69 @@ class FiniteDifferenceStep(MessagePassing):
TODO: add docstring.
"""
def __init__(self, aggr: str = "add", root_weight: float = 1.0):
def __init__(self, edge_ch=5, hidden_dim=16, aggr: str = "add"):
super().__init__(aggr=aggr)
assert (
aggr == "add"
), "Per somme pesate, l'aggregazione deve essere 'add'."
self.correction_net = nn.Sequential(
nn.Linear(2, 6),
nn.Tanh(),
nn.Linear(6, 1),
nn.Tanh(),
self.x_embedding = nn.Sequential(
spectral_norm(nn.Linear(1, hidden_dim // 2)),
nn.GELU(),
spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
)
self.edge_embedding = nn.Sequential(
spectral_norm(nn.Linear(edge_ch, hidden_dim // 2)),
nn.GELU(),
spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
)
self.update_net = nn.Sequential(
spectral_norm(nn.Linear(1, 6)),
nn.Softplus(),
spectral_norm(nn.Linear(6, 1)),
nn.Softplus(),
spectral_norm(nn.Linear(2 * hidden_dim, hidden_dim)),
nn.GELU(),
spectral_norm(nn.Linear(hidden_dim, hidden_dim)),
nn.GELU(),
# spectral_norm(nn.Linear(hidden_dim // 2, 1)),
)
self.message_net = nn.Sequential(
spectral_norm(nn.Linear(1, 6)),
nn.Softplus(),
spectral_norm(nn.Linear(6, 1)),
nn.Softplus(),
# self.message_net = nn.Sequential(
# spectral_norm(nn.Linear(2 * hidden_dim, hidden_dim)),
# nn.GELU(),
# spectral_norm(nn.Linear(hidden_dim, 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.p = torch.nn.Parameter(torch.tensor(0.5))
# self.a = torch.nn.Parameter(torch.tensor(root_weight))
def forward(self, x, edge_index, edge_attr, deg):
"""
TODO: add docstring.
"""
out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg)
return out
x_ = self.x_embedding(x)
edge_attr_ = self.edge_embedding(edge_attr)
out = self.propagate(edge_index, x=x_, edge_attr=edge_attr_, deg=deg)
return self.out_net(x_ + out)
def message(self, x_j, edge_attr):
"""
TODO: add docstring.
"""
# x_in = torch.cat([x_j, edge_attr.view(-1, 1)], dim=-1)
# correction = self.correction_net(x_in)
# p = torch.sigmoid(self.p)
# return (p * edge_attr.view(-1, 1) + (1 - p) * correction) * x_j
return edge_attr.view(-1, 1) * x_j
# msg_input = torch.cat([x_j, edge_attr], dim=-1)
# return self.message_net(msg_input) * edge_attr[:, 3].view(-1, 1)
return x_j * edge_attr
def update(self, aggr_out, x):
"""
TODO: add docstring.
"""
update_input = torch.cat([x, aggr_out], dim=-1)
return self.update_net(update_input)
# return self.update_net(aggr_out)
# return aggr_out
# h = self.update_net(aggr_out, x)
# return h
def aggregate(self, inputs, index, deg):
"""
@@ -62,68 +115,12 @@ class FiniteDifferenceStep(MessagePassing):
deg = deg + 1e-7
return out / deg.view(-1, 1)
def update(self, aggr_out, x):
"""
TODO: add docstring.
"""
return self.update_net(aggr_out)
# # Da fare:
# # - Finire calcolo della loss su ogni step e poi media
# # - Test con vari modelli
# # - Se non dovesse funzionare, provare ad adeguare il criterio di uscita
class GraphFiniteDifference(nn.Module):
"""
TODO: add docstring.
"""
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", root_weight=1.0)
@staticmethod
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_attr)
return deg + 1e-7
@staticmethod
def _compute_c_ij(c, edge_index):
"""
TODO: add docstring.
"""
return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze()
def forward(
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_attr = edge_attr * c_ij
deg = self._compute_deg(edge_index, edge_attr, x.size(0))
conv_thres = self.threshold * torch.norm(x.detach())
for _i in range(self.max_iters):
out = self.fd_step(x, edge_index, edge_attr, deg)
out[boundary_mask] = boundary_values.unsqueeze(-1)
with torch.no_grad():
residual_norm = torch.norm(out - x)
if residual_norm < conv_thres:
break
x = out.detach()
return out, _i + 1
# # PINN batching:
# # - Provare singola condizione
# # - Ottimizzatore del secondo ordine (LBFGS)