286 lines
8.9 KiB
Python
286 lines
8.9 KiB
Python
import torch
|
|
from lightning import LightningModule
|
|
from torch_geometric.data import Batch
|
|
import importlib
|
|
from matplotlib import pyplot as plt
|
|
from matplotlib.tri import Triangulation
|
|
from .model.finite_difference import FiniteDifferenceStep
|
|
import os
|
|
|
|
|
|
def import_class(class_path: str):
|
|
module_path, class_name = class_path.rsplit(".", 1) # split last dot
|
|
module = importlib.import_module(module_path) # import the module
|
|
cls = getattr(module, class_name) # get the class
|
|
return cls
|
|
|
|
|
|
def _plot_mesh(pos, y, y_pred, batch, i, batch_idx):
|
|
|
|
idx = batch == 0
|
|
y = y[idx].detach().cpu()
|
|
y_pred = y_pred[idx].detach().cpu()
|
|
pos = pos[idx].detach().cpu()
|
|
folder = f"{batch_idx:02d}_images"
|
|
if os.path.exists(folder) is False:
|
|
os.makedirs(folder)
|
|
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"{folder}/graph_iter_{i:04d}.png"
|
|
plt.savefig(name, dpi=72)
|
|
plt.close()
|
|
|
|
|
|
def _plot_losses(losses, batch_idx):
|
|
folder = f"{batch_idx:02d}_images"
|
|
plt.figure()
|
|
plt.plot(losses)
|
|
plt.yscale("log")
|
|
plt.xlabel("Iteration")
|
|
plt.ylabel("Loss")
|
|
plt.title("Test Loss over Iterations")
|
|
plt.grid(True)
|
|
file_name = f"{folder}/test_loss.png"
|
|
plt.savefig(file_name, dpi=300)
|
|
plt.close()
|
|
|
|
|
|
class GraphSolver(LightningModule):
|
|
def __init__(
|
|
self,
|
|
model_class_path: str,
|
|
model_init_args: dict = {},
|
|
loss: torch.nn.Module = None,
|
|
curriculum_learning: bool = False,
|
|
start_iters: int = 10,
|
|
increase_every: int = 100,
|
|
increase_rate: float = 1.1,
|
|
max_iters: int = 1000,
|
|
accumulation_iters: int = None,
|
|
):
|
|
super().__init__()
|
|
self.model = import_class(model_class_path)(**model_init_args)
|
|
self.fd_net = FiniteDifferenceStep()
|
|
self.loss = loss if loss is not None else torch.nn.MSELoss()
|
|
self.curriculum_learning = curriculum_learning
|
|
self.start_iters = start_iters
|
|
self.increase_every = increase_every
|
|
self.increase_rate = increase_rate
|
|
self.max_iters = max_iters
|
|
self.current_iters = start_iters
|
|
self.accumulation_iters = accumulation_iters
|
|
self.automatic_optimization = False
|
|
self.threshold = 1e-5
|
|
|
|
self.alpha = torch.nn.Parameter(torch.tensor(0.1))
|
|
|
|
def _compute_deg(self, edge_index, edge_attr, num_nodes):
|
|
deg = torch.zeros(num_nodes, device=edge_index.device)
|
|
deg = deg.scatter_add(0, edge_index[1], edge_attr)
|
|
return deg + 1e-7
|
|
|
|
def _compute_loss(self, x, y):
|
|
return self.loss(x, y)
|
|
|
|
def _log_loss(self, loss, batch, stage: str):
|
|
self.log(
|
|
f"{stage}/loss",
|
|
loss,
|
|
on_step=False,
|
|
on_epoch=True,
|
|
prog_bar=True,
|
|
batch_size=int(batch.num_graphs),
|
|
)
|
|
return loss
|
|
|
|
@staticmethod
|
|
def _compute_c_ij(c, edge_index):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze()
|
|
|
|
def _compute_model_steps(
|
|
self, x, edge_index, edge_attr, deg, boundary_mask, boundary_values
|
|
):
|
|
# with torch.no_grad():
|
|
# out = self.fd_net(x, edge_index, edge_attr, deg)
|
|
# out[boundary_mask] = boundary_values.unsqueeze(-1)
|
|
# diff = out - x
|
|
# out = self.model(out, edge_index, edge_attr, deg)
|
|
# out = out + self.alpha * correction
|
|
# out[boundary_mask] = boundary_values.unsqueeze(-1)
|
|
out = self.model(x, edge_index, edge_attr, deg)
|
|
out[boundary_mask] = boundary_values.unsqueeze(-1)
|
|
return out
|
|
|
|
def _check_convergence(self, out, x):
|
|
residual_norm = torch.norm(out - x)
|
|
if residual_norm < self.threshold * torch.norm(x):
|
|
return True
|
|
return False
|
|
|
|
def accumulate_gradients(self, losses):
|
|
loss_ = torch.stack(losses, dim=0).mean()
|
|
self.manual_backward(loss_, retain_graph=True)
|
|
return loss_.item()
|
|
|
|
def _preprocess_batch(self, batch: Batch):
|
|
x, y, c, edge_index, edge_attr = (
|
|
batch.x,
|
|
batch.y,
|
|
batch.c,
|
|
batch.edge_index,
|
|
batch.edge_attr,
|
|
)
|
|
edge_attr = 1 / edge_attr
|
|
c_ij = self._compute_c_ij(c, edge_index)
|
|
edge_attr = edge_attr * c_ij
|
|
return x, y, edge_index, edge_attr
|
|
|
|
def training_step(self, batch: Batch, _):
|
|
optim = self.optimizers()
|
|
optim.zero_grad()
|
|
x, y, edge_index, edge_attr = self._preprocess_batch(batch)
|
|
|
|
deg = self._compute_deg(edge_index, edge_attr, x.size(0))
|
|
losses = []
|
|
acc_loss, acc_it = 0, 0
|
|
|
|
for i in range(self.current_iters):
|
|
out = self._compute_model_steps(
|
|
x,
|
|
edge_index,
|
|
edge_attr.unsqueeze(-1),
|
|
deg,
|
|
batch.boundary_mask,
|
|
batch.boundary_values,
|
|
)
|
|
losses.append(self.loss(out, y))
|
|
|
|
# Accumulate gradients if reached accumulation iters
|
|
if (
|
|
self.accumulation_iters is not None
|
|
and (i + 1) % self.accumulation_iters == 0
|
|
):
|
|
loss = self.accumulate_gradients(losses)
|
|
losses = []
|
|
acc_it += 1
|
|
out = out.detach()
|
|
acc_loss = acc_loss + loss
|
|
|
|
# Check for convergence and break if converged (with final accumulation)
|
|
converged = self._check_convergence(out, x)
|
|
if converged:
|
|
if losses:
|
|
loss = self.accumulate_gradients(losses)
|
|
acc_it += 1
|
|
acc_loss = acc_loss + loss
|
|
break
|
|
|
|
# Final accumulation if we are at the last iteration
|
|
if i == self.current_iters - 1:
|
|
if losses:
|
|
loss = self.accumulate_gradients(losses)
|
|
acc_it += 1
|
|
acc_loss = acc_loss + loss
|
|
|
|
x = out
|
|
|
|
loss = self.loss(out, y)
|
|
for param in self.model.parameters():
|
|
if param.grad is not None:
|
|
param.grad /= acc_it
|
|
optim.step()
|
|
optim.zero_grad()
|
|
|
|
self.log(
|
|
"train/accumulated_loss",
|
|
(acc_loss / acc_it if acc_it > 0 else acc_loss),
|
|
on_step=False,
|
|
on_epoch=True,
|
|
prog_bar=True,
|
|
batch_size=int(batch.num_graphs),
|
|
)
|
|
self.log(
|
|
"train/iterations",
|
|
i + 1,
|
|
on_step=False,
|
|
on_epoch=True,
|
|
prog_bar=True,
|
|
batch_size=int(batch.num_graphs),
|
|
)
|
|
if hasattr(self.model, "p"):
|
|
self.log(
|
|
"train/p",
|
|
self.model.p,
|
|
on_step=False,
|
|
on_epoch=True,
|
|
prog_bar=True,
|
|
batch_size=int(batch.num_graphs),
|
|
)
|
|
|
|
def on_train_epoch_end(self):
|
|
if self.curriculum_learning:
|
|
if (self.current_iters < self.max_iters) and (
|
|
self.current_epoch % self.increase_every == 0
|
|
):
|
|
self.current_iters = min(
|
|
int(self.current_iters * self.increase_rate), self.max_iters
|
|
)
|
|
return super().on_train_epoch_end()
|
|
|
|
def validation_step(self, batch: Batch, _):
|
|
x, y, edge_index, edge_attr = self._preprocess_batch(batch)
|
|
|
|
deg = self._compute_deg(edge_index, edge_attr, x.size(0))
|
|
for i in range(self.current_iters):
|
|
out = self._compute_model_steps(
|
|
x,
|
|
edge_index,
|
|
edge_attr.unsqueeze(-1),
|
|
deg,
|
|
batch.boundary_mask,
|
|
batch.boundary_values,
|
|
)
|
|
converged = self._check_convergence(out, x)
|
|
|
|
if converged:
|
|
break
|
|
x = out
|
|
loss = self.loss(out, y)
|
|
self._log_loss(loss, batch, "val")
|
|
self.log(
|
|
"val/iterations",
|
|
i + 1,
|
|
on_step=False,
|
|
on_epoch=True,
|
|
prog_bar=True,
|
|
batch_size=int(batch.num_graphs),
|
|
)
|
|
|
|
def test_step(self, batch: Batch, batch_idx):
|
|
pass
|
|
|
|
def configure_optimizers(self):
|
|
optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3)
|
|
return optimizer
|
|
|
|
def _impose_bc(self, x: torch.Tensor, data: Batch):
|
|
x[data.boundary_mask] = data.boundary_values
|
|
return x
|