Skip to content

Fix device mismatch for residual weights in RBAPINN #614

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dario-coscia merged 1 commit into from
Jul 31, 2025
Merged

Fix device mismatch for residual weights in RBAPINN #614

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 18 additions & 3 deletions pina/solver/physics_informed_solver/rba_pinn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,21 @@ def __init__(
# Set the loss function to return non-aggregated losses
self._loss_fn = type(self._loss_fn)(reduction="none")

def on_train_start(self):
"""
Ensure that all residual weight buffers registered during initialization
are moved to the correct computation device.
"""
# Move all weight buffers to the correct device
for cond in self.problem.input_pts:

# Get the buffer for the current condition
weight_buf = getattr(self, f"weight_{cond}")

# Move the buffer to the correct device
weight_buf.data = weight_buf.data.to(self.device)
self.weights[cond] = weight_buf

def training_step(self, batch, batch_idx, **kwargs):
"""
Solver training step. It computes the optimization cycle and aggregates
Expand Down Expand Up @@ -235,7 +250,7 @@ def _optimization_cycle(self, batch, batch_idx, **kwargs):
idx = torch.arange(
batch_idx * len_res,
(batch_idx + 1) * len_res,
device=res.device,
device=self.weights[cond].device,
) % len(self.problem.input_pts[cond])

losses[cond] = self._apply_reduction(
Expand Down Expand Up @@ -271,7 +286,7 @@ def _update_weights(self, batch, batch_idx, residuals):

# Compute normalized residuals
res = residuals[cond]
res_abs = res.abs()
res_abs = torch.linalg.vector_norm(res, ord=2, dim=1, keepdim=True)

@dario-coscia dario-coscia Jul 30, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is it changing from absolute value to norm l2?

@GiovanniCanali GiovanniCanali Jul 30, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is done to ensure that the RBAPINN solver works also for a condition associated with a SystemEquation. Indeed, the weights' update must be a scalar.

Notice that when dealing with a single Equation, the L2 norm along the first dimension is the same as the absolute value.

@GiovanniCanali GiovanniCanali Jul 30, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@dario-coscia

r_norm = (self.eta * res_abs) / (res_abs.max() + 1e-12)

# Get the correct indices for the weights. Modulus is used according
Expand All @@ -280,7 +295,7 @@ def _update_weights(self, batch, batch_idx, residuals):
idx = torch.arange(
batch_idx * len_pts,
(batch_idx + 1) * len_pts,
device=res.device,
device=self.weights[cond].device,
) % len(self.problem.input_pts[cond])

# Update weights
Expand Down
Loading