NVIDIA Modulus Sym v1.3.0
v1.3.0

deeplearning/modulus/modulus-sym-v130/_modules/modulus/sym/loss/aggregator.html

Source code for modulus.sym.loss.aggregator

# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Import libraries
import torch
import logging
import numpy as np
from torch import nn
from typing import Dict, List, Optional, Callable, Union

# Import from Modulus
from modulus.sym.eq.derivatives import gradient
from modulus.sym.hydra import to_absolute_path, add_hydra_run_path

logger = logging.getLogger(__name__)


[docs]class Aggregator(nn.Module): """ Base class for loss aggregators """ def __init__(self, params, num_losses, weights): super().__init__() self.params: List[torch.Tensor] = list(params) self.num_losses: int = num_losses self.weights: Optional[Dict[str, float]] = weights self.device: torch.device self.device = list(set(p.device for p in self.params))[0] self.init_loss: torch.Tensor = torch.tensor(0.0, device=self.device) def weigh_losses_initialize( weights: Optional[Dict[str, float]] ) -> Callable[ [Dict[str, torch.Tensor], Optional[Dict[str, float]]], Dict[str, torch.Tensor], ]: if weights is None: def weigh_losses( losses: Dict[str, torch.Tensor], weights: None ) -> Dict[str, torch.Tensor]: return losses else: def weigh_losses( losses: Dict[str, torch.Tensor], weights: Dict[str, float] ) -> Dict[str, torch.Tensor]: for key in losses.keys(): if key not in weights.keys(): weights.update({key: 1.0}) losses = {key: weights[key] * losses[key] for key in losses.keys()} return losses return weigh_losses self.weigh_losses = weigh_losses_initialize(self.weights)
[docs]class Sum(Aggregator): """ Loss aggregation by summation """ def __init__(self, params, num_losses, weights=None): super().__init__(params, num_losses, weights)
[docs] def forward(self, losses: Dict[str, torch.Tensor], step: int) -> torch.Tensor: """ Aggregates the losses by summation Parameters ---------- losses : Dict[str, torch.Tensor] A dictionary of losses. step : int Optimizer step. Returns ------- loss : torch.Tensor Aggregated loss. """ # weigh losses losses = self.weigh_losses(losses, self.weights) # Initialize loss loss: torch.Tensor = torch.zeros_like(self.init_loss) # Add losses for key in losses.keys(): loss += losses[key] return loss
[docs]class GradNorm(Aggregator): """ GradNorm for loss aggregation Reference: "Chen, Z., Badrinarayanan, V., Lee, C.Y. and Rabinovich, A., 2018, July. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning (pp. 794-803). PMLR." """ def __init__(self, params, num_losses, alpha=1.0, weights=None): super().__init__(params, num_losses, weights) self.alpha: float = alpha self.lmbda: torch.nn.Parameter = nn.Parameter( torch.zeros(num_losses, device=self.device) ) self.register_buffer( "init_losses", torch.zeros(self.num_losses, device=self.device) )
[docs] def forward(self, losses: Dict[str, torch.Tensor], step: int) -> torch.Tensor: """ Weights and aggregates the losses using the gradNorm algorithm Parameters ---------- losses : Dict[str, torch.Tensor] A dictionary of losses. step : int Optimizer step. Returns ------- loss : torch.Tensor Aggregated loss. """ # weigh losses losses = self.weigh_losses(losses, self.weights) # get initial losses if step == 0: for i, key in enumerate(losses.keys()): self.init_losses[i] = losses[key].clone().detach() with torch.no_grad(): normalizer: torch.Tensor = self.num_losses / (torch.exp(self.lmbda).sum()) for i in range(self.num_losses): self.lmbda[i] = self.lmbda[i].clone() + torch.log( normalizer.detach() ) # c*exp(x) = exp(log(c)+x) lmbda_exp: torch.Tensor = torch.exp(self.lmbda) # compute relative losses, inverse rate, and grad coefficient losses_stacked: torch.Tensor = torch.stack(list(losses.values())) with torch.no_grad(): relative_losses: torch.Tensor = torch.div(losses_stacked, self.init_losses) inverse_rate: torch.Tensor = relative_losses / (relative_losses.mean()) gradnorm_coef: torch.Tensor = torch.pow(inverse_rate, self.alpha) # compute gradient norm and average gradient norm grads_norm: torch.Tensor = torch.zeros_like(self.init_losses) shared_params: torch.Tensor = self.params[-2] # TODO generalize this for i, key in enumerate(losses.keys()): grads: torch.Tensor = gradient(losses[key], [shared_params])[0] grads_norm[i] = torch.norm(lmbda_exp[i] * grads.detach(), p=2) avg_grad: torch.Tensor = grads_norm.detach().mean() # compute gradnorm & model losses loss_gradnorm: torch.Tensor = torch.abs( grads_norm - avg_grad * gradnorm_coef ).sum() loss_model: torch.Tensor = (lmbda_exp.detach() * losses_stacked).sum() loss: torch.Tensor = loss_gradnorm + loss_model return loss
[docs]class ResNorm(Aggregator): """ Residual normalization for loss aggregation Contributors: T. Nandi, D. Van Essendelft, M. A. Nabian """ def __init__(self, params, num_losses, alpha=1.0, weights=None): super().__init__(params, num_losses, weights) self.alpha: float = alpha self.lmbda: torch.nn.Parameter = nn.Parameter( torch.zeros(num_losses, device=self.device) ) self.register_buffer( "init_losses", torch.zeros(self.num_losses, device=self.device) )
[docs] def forward(self, losses: Dict[str, torch.Tensor], step: int) -> torch.Tensor: """ Weights and aggregates the losses using the ResNorm algorithm Parameters ---------- losses : Dict[str, torch.Tensor] A dictionary of losses. step : int Optimizer step. Returns ------- loss : torch.Tensor Aggregated loss. """ # weigh losses losses = self.weigh_losses(losses, self.weights) # get initial losses if step == 0: for i, key in enumerate(losses.keys()): self.init_losses[i] = losses[key].clone().detach() with torch.no_grad(): normalizer: torch.Tensor = self.num_losses / (torch.exp(self.lmbda).sum()) for i in range(self.num_losses): self.lmbda[i] = self.lmbda[i].clone() + torch.log( normalizer.detach() ) # c*exp(x) = exp(log(c)+x) lmbda_exp: torch.Tensor = torch.exp(self.lmbda) # compute relative losses, inverse rate, and grad coefficient losses_stacked: torch.Tensor = torch.stack(list(losses.values())) with torch.no_grad(): relative_losses: torch.Tensor = torch.div(losses_stacked, self.init_losses) inverse_rate: torch.Tensor = relative_losses / (relative_losses.mean()) resnorm_coef: torch.Tensor = torch.pow(inverse_rate, self.alpha) # compute residual norm and average residual norm residuals: torch.Tensor = torch.zeros_like(self.init_losses) for i, key in enumerate(losses.keys()): residuals[i] = lmbda_exp[i] * losses[key].detach() avg_residuals: torch.Tensor = losses_stacked.detach().mean() # compute ResNorm & model losses loss_resnorm: torch.Tensor = torch.abs( residuals - avg_residuals * resnorm_coef ).sum() loss_model: torch.Tensor = (lmbda_exp.detach() * losses_stacked).sum() loss: torch.Tensor = loss_resnorm + loss_model return loss
[docs]class HomoscedasticUncertainty(Aggregator): """ Homoscedastic task uncertainty for loss aggregation Reference: "Reference: Kendall, A., Gal, Y. and Cipolla, R., 2018. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 7482-7491)." """ def __init__(self, params, num_losses, weights=None): super().__init__(params, num_losses, weights) self.log_var: torch.nn.Parameter = nn.Parameter( torch.zeros(self.num_losses, device=self.device) )
[docs] def forward(self, losses: Dict[str, torch.Tensor], step: int) -> torch.Tensor: """ Weights and aggregates the losses using homoscedastic task uncertainty Parameters ---------- losses : Dict[str, torch.Tensor] A dictionary of losses. step : int Optimizer step. Returns ------- loss : torch.Tensor Aggregated loss. """ # weigh losses losses = self.weigh_losses(losses, self.weights) # Initialize loss loss: torch.Tensor = torch.zeros_like(self.init_loss) # Compute precision precision: torch.Tensor = torch.exp(-self.log_var) # Aggregate losses for i, key in enumerate(losses.keys()): loss += precision[i] * losses[key] loss += self.log_var.sum() loss /= 2.0 return loss
[docs]class LRAnnealing(Aggregator): """ Learning rate annealing for loss aggregation References: "Wang, S., Teng, Y. and Perdikaris, P., 2020. Understanding and mitigating gradient pathologies in physics-informed neural networks. arXiv preprint arXiv:2001.04536.", and "Jin, X., Cai, S., Li, H. and Karniadakis, G.E., 2021. NSFnets (Navier-Stokes flow nets): Physics-informed neural networks for the incompressible Navier-Stokes equations. Journal of Computational Physics, 426, p.109951." """ def __init__( self, params, num_losses, update_freq=1, alpha=0.01, ref_key=None, eps=1e-8, weights=None, ): super().__init__(params, num_losses, weights) self.update_freq: int = update_freq self.alpha: float = alpha self.ref_key: Union[str, None] = ref_key self.eps: float = eps self.register_buffer( "lmbda_ema", torch.ones(self.num_losses, device=self.device) )
[docs] def forward(self, losses: Dict[str, torch.Tensor], step: int) -> torch.Tensor: """ Weights and aggregates the losses using the learning rate annealing algorithm Parameters ---------- losses : Dict[str, torch.Tensor] A dictionary of losses. step : int Optimizer step. Returns ------- loss : torch.Tensor Aggregated loss. """ # weigh losses losses = self.weigh_losses(losses, self.weights) # Initialize loss loss: torch.Tensor = torch.zeros_like(self.init_loss) # Determine reference loss if self.ref_key is None: ref_idx = 0 else: for i, key in enumerate(losses.keys()): if self.ref_key in key: ref_idx = i break # Update loss weights and aggregate losses if step % self.update_freq == 0: grads_mean: List[torch.Tensor] = [] # Compute the mean of each loss gradients for key in losses.keys(): grads: List[torch.Tensor] = gradient(losses[key], self.params) grads_flattened: List[torch.Tensor] = [] for i in range(len(grads)): if grads[i] is not None: grads_flattened.append(torch.abs(torch.flatten(grads[i]))) grads_mean.append((torch.mean(torch.cat(grads_flattened)))) # Compute the exponential moving average of weights and aggregate losses for i, key in enumerate(losses.keys()): with torch.no_grad(): self.lmbda_ema[i] *= 1.0 - self.alpha self.lmbda_ema[i] += ( self.alpha * grads_mean[ref_idx] / (grads_mean[i] + self.eps) ) loss += self.lmbda_ema[i].clone() * losses[key] # Aggregate losses without update to loss weights else: for i, key in enumerate(losses.keys()): loss += self.lmbda_ema[i] * losses[key] return loss
[docs]class SoftAdapt(Aggregator): """ SoftAdapt for loss aggregation Reference: "Heydari, A.A., Thompson, C.A. and Mehmood, A., 2019. Softadapt: Techniques for adaptive loss weighting of neural networks with multi-part loss functions. arXiv preprint arXiv: 1912.12355." """ def __init__(self, params, num_losses, eps=1e-8, weights=None): super().__init__(params, num_losses, weights) self.eps: float = eps self.register_buffer( "prev_losses", torch.zeros(self.num_losses, device=self.device) )
[docs] def forward(self, losses: Dict[str, torch.Tensor], step: int) -> torch.Tensor: """ Weights and aggregates the losses using the original variant of the softadapt algorithm Parameters ---------- losses : Dict[str, torch.Tensor] A dictionary of losses. step : int Optimizer step. Returns ------- loss : torch.Tensor Aggregated loss. """ # weigh losses losses = self.weigh_losses(losses, self.weights) # Initialize loss loss: torch.Tensor = torch.zeros_like(self.init_loss) # Aggregate losses by summation at step 0 if step == 0: for i, key in enumerate(losses.keys()): loss += losses[key] self.prev_losses[i] = losses[key].clone().detach() # Aggregate losses using SoftAdapt for step > 0 else: lmbda: torch.Tensor = torch.ones_like(self.prev_losses) lmbda_sum: torch.Tensor = torch.zeros_like(self.init_loss) losses_stacked: torch.Tensor = torch.stack(list(losses.values())) normalizer: torch.Tensor = (losses_stacked / self.prev_losses).max() for i, key in enumerate(losses.keys()): with torch.no_grad(): lmbda[i] = torch.exp( losses[key] / (self.prev_losses[i] + self.eps) - normalizer ) lmbda_sum += lmbda[i] loss += lmbda[i].clone() * losses[key] self.prev_losses[i] = losses[key].clone().detach() loss *= self.num_losses / (lmbda_sum + self.eps) return loss
[docs]class Relobralo(Aggregator): """ Relative loss balancing with random lookback Reference: "Bischof, R. and Kraus, M., 2021. Multi-Objective Loss Balancing for Physics-Informed Deep Learning. arXiv preprint arXiv:2110.09813." """ def __init__( self, params, num_losses, alpha=0.95, beta=0.99, tau=1.0, eps=1e-8, weights=None ): super().__init__(params, num_losses, weights) self.alpha: float = alpha self.beta: float = beta self.tau: float = tau self.eps: float = eps self.register_buffer( "init_losses", torch.zeros(self.num_losses, device=self.device) ) self.register_buffer( "prev_losses", torch.zeros(self.num_losses, device=self.device) ) self.register_buffer( "lmbda_ema", torch.ones(self.num_losses, device=self.device) )
[docs] def forward(self, losses: Dict[str, torch.Tensor], step: int) -> torch.Tensor: """ Weights and aggregates the losses using the ReLoBRaLo algorithm Parameters ---------- losses : Dict[str, torch.Tensor] A dictionary of losses. step : int Optimizer step. Returns ------- loss : torch.Tensor Aggregated loss. """ # weigh losses losses = self.weigh_losses(losses, self.weights) # Initialize loss loss: torch.Tensor = torch.zeros_like(self.init_loss) # Aggregate losses by summation at step 0 if step == 0: for i, key in enumerate(losses.keys()): loss += losses[key] self.init_losses[i] = losses[key].clone().detach() self.prev_losses[i] = losses[key].clone().detach() # Aggregate losses using ReLoBRaLo for step > 0 else: losses_stacked: torch.Tensor = torch.stack(list(losses.values())) normalizer_prev: torch.Tensor = ( losses_stacked / (self.tau * self.prev_losses) ).max() normalizer_init: torch.Tensor = ( losses_stacked / (self.tau * self.init_losses) ).max() rho: torch.Tensor = torch.bernoulli(torch.tensor(self.beta)) with torch.no_grad(): lmbda_prev: torch.Tensor = torch.exp( losses_stacked / (self.tau * self.prev_losses + self.eps) - normalizer_prev ) lmbda_init: torch.Tensor = torch.exp( losses_stacked / (self.tau * self.init_losses + self.eps) - normalizer_init ) lmbda_prev *= self.num_losses / (lmbda_prev.sum() + self.eps) lmbda_init *= self.num_losses / (lmbda_init.sum() + self.eps) # Compute the exponential moving average of weights and aggregate losses for i, key in enumerate(losses.keys()): with torch.no_grad(): self.lmbda_ema[i] = self.alpha * ( rho * self.lmbda_ema[i].clone() + (1.0 - rho) * lmbda_init[i] ) self.lmbda_ema[i] += (1.0 - self.alpha) * lmbda_prev[i] loss += self.lmbda_ema[i].clone() * losses[key] self.prev_losses[i] = losses[key].clone().detach() return loss
[docs]class NTK(nn.Module): def __init__(self, run_per_step: int = 1000, save_name: Union[str, None] = None): super(NTK, self).__init__() self.run_per_step = run_per_step self.if_csv_head = True self.save_name = ( to_absolute_path(add_hydra_run_path(save_name)) if save_name else None ) if self.save_name: logger.warning( "Cuda graphs does not work when saving NTK values to file! Set `cuda_graphs` to false." ) def group_ntk(self, model, losses): # The item in this losses should scalar loss values after MSE, etc. ntk_value = dict() for key, loss in losses.items(): grad = torch.autograd.grad( torch.sqrt(torch.abs(loss)), model.parameters(), retain_graph=True, allow_unused=True, ) ntk_value[key] = torch.sqrt( torch.sum( torch.stack( [torch.sum(t.detach() ** 2) for t in grad if t is not None], dim=0, ) ) ) return ntk_value def save_ntk(self, ntk_dict, step): import pandas as pd # TODO: Remove output_dict = {} for key, value in ntk_dict.items(): output_dict[key] = value.cpu().numpy() df = pd.DataFrame(output_dict, index=[step]) df.to_csv(self.save_name + ".csv", mode="a", header=self.if_csv_head) self.if_csv_head = False
[docs] def forward(self, constraints, ntk_weights, step): losses = dict() dict_constraint_losses = dict() ntk_sum = 0 # Execute constraint forward passes for key, constraint in constraints.items(): # TODO: Test streaming here torch.cuda.nvtx.range_push(f"Running Constraint {key}") constraint.forward() torch.cuda.nvtx.range_pop() for key, constraint in constraints.items(): # compute losses constraint_losses = constraint.loss(step) if (step % self.run_per_step == 0) and (step > 0): ntk_dict = self.group_ntk(constraint.model, constraint_losses) else: ntk_dict = None if ntk_dict is not None: ntk_weights[key] = ntk_dict if ntk_weights.get(key) is not None: ntk_sum += torch.sum( torch.stack(list(ntk_weights[key].values()), dim=0) ) dict_constraint_losses[key] = constraint_losses if step == 0: # May not work on restarts ntk_sum = 1.0 if self.save_name and (step % self.run_per_step == 0) and (step > 0): self.save_ntk( { d_key + "_" + k: v for d_key, d in ntk_weights.items() for k, v in d.items() }, step, ) for key, constraint_losses in dict_constraint_losses.items(): # add together losses of like kind for loss_key, value in constraint_losses.items(): if ( ntk_weights.get(key) is None or ntk_weights[key].get(loss_key) is None ): ntk_weight = ntk_sum / 1.0 else: ntk_weight = ntk_sum / ntk_weights[key][loss_key] if loss_key not in list(losses.keys()): losses[loss_key] = ntk_weight * value else: losses[loss_key] += ntk_weight * value return losses, ntk_weights
© Copyright 2023, NVIDIA Modulus Team. Last updated on Jan 25, 2024.