deeplearning/modulus/modulus-v2209/_modules/modulus/domain/validator/discrete.html

v22.09

Source code for modulus.domain.validator.discrete

from typing import Dict, List

import torch
import numpy as np

from modulus.domain.validator import Validator
from modulus.domain.constraint import Constraint
from modulus.utils.io.vtk import grid_to_vtk
from modulus.utils.io import GridValidatorPlotter, DeepONetValidatorPlotter
from modulus.graph import Graph
from modulus.key import Key
from modulus.node import Node
from modulus.constants import TF_SUMMARY
from modulus.distributed import DistributedManager
from modulus.dataset import Dataset, DictGridDataset


[docs]class GridValidator(Validator): """Data-driven grid field validator Parameters ---------- nodes : List[Node] List of Modulus Nodes to unroll graph with. dataset: Dataset dataset which contains invar and true outvar examples batch_size : int, optional Batch size used when running validation, by default 100 plotter : GridValidatorPlotter Modulus plotter for showing results in tensorboard. requires_grad : bool = False If automatic differentiation is needed for computing results. num_workers : int, optional Number of dataloader workers, by default 0 """ def __init__( self, nodes: List[Node], dataset: Dataset, batch_size: int = 100, plotter: GridValidatorPlotter = None, requires_grad: bool = False, num_workers: int = 0, ): # get dataset and dataloader self.dataset = dataset self.dataloader = Constraint.get_dataloader( dataset=self.dataset, batch_size=batch_size, shuffle=False, drop_last=False, num_workers=num_workers, distributed=False, infinite=False, ) # construct model from nodes self.model = Graph( nodes, Key.convert_list(self.dataset.invar_keys), Key.convert_list(self.dataset.outvar_keys), ) self.manager = DistributedManager() self.device = self.manager.device self.model.to(self.device) # set foward method self.requires_grad = requires_grad self.forward = self.forward_grad if requires_grad else self.forward_nograd # set plotter self.plotter = plotter def save_results(self, name, results_dir, writer, save_filetypes, step): invar_cpu = {key: [] for key in self.dataset.invar_keys} true_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} pred_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} # Loop through mini-batches for i, (invar0, true_outvar0, lambda_weighting) in enumerate(self.dataloader): # Move data to device (may need gradients in future, if so requires_grad=True) invar = Constraint._set_device( invar0, device=self.device, requires_grad=self.requires_grad ) true_outvar = Constraint._set_device( true_outvar0, device=self.device, requires_grad=self.requires_grad ) pred_outvar = self.forward(invar) # Collect minibatch info into cpu dictionaries invar_cpu = { key: value + [invar[key].cpu().detach()] for key, value in invar_cpu.items() } true_outvar_cpu = { key: value + [true_outvar[key].cpu().detach()] for key, value in true_outvar_cpu.items() } pred_outvar_cpu = { key: value + [pred_outvar[key].cpu().detach()] for key, value in pred_outvar_cpu.items() } # Concat mini-batch tensors invar_cpu = {key: torch.cat(value) for key, value in invar_cpu.items()} true_outvar_cpu = { key: torch.cat(value) for key, value in true_outvar_cpu.items() } pred_outvar_cpu = { key: torch.cat(value) for key, value in pred_outvar_cpu.items() } # compute losses on cpu losses = GridValidator._l2_relative_error(true_outvar_cpu, pred_outvar_cpu) # convert to numpy arrays invar = {k: v.numpy() for k, v in invar_cpu.items()} true_outvar = {k: v.numpy() for k, v in true_outvar_cpu.items()} pred_outvar = {k: v.numpy() for k, v in pred_outvar_cpu.items()} # save batch to vtk file TODO clean this up after graph unroll stuff named_true_outvar = {"true_" + k: v for k, v in true_outvar.items()} named_pred_outvar = {"pred_" + k: v for k, v in pred_outvar.items()} # save batch to vtk/npz file TODO clean this up after graph unroll stuff if "np" in save_filetypes: np.savez( results_dir + name, {**invar, **named_true_outvar, **named_pred_outvar} ) if "vtk" in save_filetypes: grid_to_vtk( {**invar, **named_true_outvar, **named_pred_outvar}, results_dir + name ) # add tensorboard plots if self.plotter is not None: self.plotter._add_figures( "Validators", name, results_dir, writer, step, invar, true_outvar, pred_outvar, ) # add tensorboard scalars for k, loss in losses.items(): if TF_SUMMARY: writer.add_scalar("val/" + name + "/" + k, loss, step, new_style=True) else: writer.add_scalar( "Validators/" + name + "/" + k, loss, step, new_style=True ) return losses

class _DeepONet_Validator(Validator): def __init__( self, nodes: List[Node], invar_branch: Dict[str, np.array], invar_trunk: Dict[str, np.array], true_outvar: Dict[str, np.array], batch_size: int, plotter: DeepONetValidatorPlotter, requires_grad: bool, ): # TODO: add support for other datasets? # get dataset and dataloader self.dataset = DictGridDataset(invar=invar_branch, outvar=true_outvar) self.dataloader = Constraint.get_dataloader( dataset=self.dataset, batch_size=batch_size, shuffle=False, drop_last=False, num_workers=0, distributed=False, infinite=False, ) # construct model from nodes self.model = Graph( nodes, Key.convert_list(invar_branch.keys()) + Key.convert_list(invar_trunk.keys()), Key.convert_list(true_outvar.keys()), ) self.manager = DistributedManager() self.device = self.manager.device self.model.to(self.device) # set foward method self.requires_grad = requires_grad self.forward = self.forward_grad if requires_grad else self.forward_nograd # set plotter self.plotter = plotter

[docs]class DeepONet_Physics_Validator(_DeepONet_Validator): """ DeepONet Validator """ def __init__( self, nodes: List[Node], invar_branch: Dict[str, np.array], invar_trunk: Dict[str, np.array], true_outvar: Dict[str, np.array], batch_size: int = 100, plotter: DeepONetValidatorPlotter = None, requires_grad: bool = False, tile_trunk_input: bool = True, ): super().__init__( nodes=nodes, invar_branch=invar_branch, invar_trunk=invar_trunk, true_outvar=true_outvar, batch_size=batch_size, plotter=plotter, requires_grad=requires_grad, ) if tile_trunk_input: for k, v in invar_trunk.items(): invar_trunk[k] = np.tile(v, (batch_size, 1)) self.invar_trunk = invar_trunk self.batch_size = batch_size def save_results(self, name, results_dir, writer, save_filetypes, step): invar_cpu = {key: [] for key in self.dataset.invar_keys} invar_trunk_gpu = Constraint._set_device( self.invar_trunk, device=self.device, requires_grad=self.requires_grad ) true_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} pred_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} # Loop through mini-batches for i, (invar0, true_outvar0, lambda_weighting) in enumerate(self.dataloader): # Move data to device (may need gradients in future, if so requires_grad=True) invar = Constraint._set_device( invar0, device=self.device, requires_grad=self.requires_grad ) true_outvar = Constraint._set_device( true_outvar0, device=self.device, requires_grad=self.requires_grad ) pred_outvar = self.forward({**invar, **invar_trunk_gpu}) # Collect minibatch info into cpu dictionaries invar_cpu = { key: value + [invar[key].cpu().detach()] for key, value in invar_cpu.items() } true_outvar_cpu = { key: value + [true_outvar[key].cpu().detach()] for key, value in true_outvar_cpu.items() } pred_outvar_cpu = { key: value + [pred_outvar[key].cpu().detach()] for key, value in pred_outvar_cpu.items() } # Concat mini-batch tensors invar_cpu = {key: torch.cat(value) for key, value in invar_cpu.items()} true_outvar_cpu = { key: torch.cat(value) for key, value in true_outvar_cpu.items() } pred_outvar_cpu = { key: torch.cat(value) for key, value in pred_outvar_cpu.items() } # compute losses on cpu losses = DeepONet_Physics_Validator._l2_relative_error( true_outvar_cpu, pred_outvar_cpu ) # convert to numpy arrays invar = {k: v.numpy() for k, v in invar_cpu.items()} true_outvar = {k: v.numpy() for k, v in true_outvar_cpu.items()} pred_outvar = {k: v.numpy() for k, v in pred_outvar_cpu.items()} # save batch to vtk file TODO clean this up after graph unroll stuff named_true_outvar = {"true_" + k: v for k, v in true_outvar.items()} named_pred_outvar = {"pred_" + k: v for k, v in pred_outvar.items()} # save batch to vtk/npz file TODO clean this up after graph unroll stuff if "np" in save_filetypes: np.savez( results_dir + name, {**invar, **named_true_outvar, **named_pred_outvar} ) ndim = next(iter(self.invar_trunk.values())).shape[-1] invar_plotter = dict() true_outvar_plotter = dict() pred_outvar_plotter = dict() for k, v in self.invar_trunk.items(): invar_plotter[k] = self.invar_trunk[k].reshape((self.batch_size, -1, ndim)) for k, v in true_outvar.items(): true_outvar_plotter[k] = true_outvar[k].reshape((self.batch_size, -1)) for k, v in pred_outvar.items(): pred_outvar_plotter[k] = pred_outvar[k].reshape((self.batch_size, -1)) # add tensorboard plots if self.plotter is not None: self.plotter._add_figures( "Validators", name, results_dir, writer, step, invar_plotter, true_outvar_plotter, pred_outvar_plotter, ) # add tensorboard scalars for k, loss in losses.items(): if TF_SUMMARY: writer.add_scalar("val/" + name + "/" + k, loss, step, new_style=True) else: writer.add_scalar( "Validators/" + name + "/" + k, loss, step, new_style=True ) return losses @staticmethod def _l2_relative_error(true_var, pred_var): # TODO replace with metric classes new_var = {} for key in true_var.keys(): new_var["l2_relative_error_" + str(key)] = torch.sqrt( torch.mean( torch.square(torch.reshape(true_var[key], (-1, 1)) - pred_var[key]) ) / torch.var(true_var[key]) ) return new_var
[docs]class DeepONet_Data_Validator(_DeepONet_Validator): """ DeepONet Validator """ def __init__( self, nodes: List[Node], invar_branch: Dict[str, np.array], invar_trunk: Dict[str, np.array], true_outvar: Dict[str, np.array], batch_size: int = 100, plotter: DeepONetValidatorPlotter = None, requires_grad: bool = False, ): super().__init__( nodes=nodes, invar_branch=invar_branch, invar_trunk=invar_trunk, true_outvar=true_outvar, batch_size=batch_size, plotter=plotter, requires_grad=requires_grad, ) self.invar_trunk_plotter = dict() ndim = next(iter(invar_trunk.values())).shape[-1] for k, v in invar_trunk.items(): self.invar_trunk_plotter[k] = np.tile(v, (batch_size, 1)).reshape( (batch_size, -1, ndim) ) self.invar_trunk = invar_trunk def save_results(self, name, results_dir, writer, save_filetypes, step): invar_cpu = {key: [] for key in self.dataset.invar_keys} invar_trunk_gpu = Constraint._set_device( self.invar_trunk, device=self.device, requires_grad=self.requires_grad ) true_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} pred_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} # Loop through mini-batches for i, (invar0, true_outvar0, lambda_weighting) in enumerate(self.dataloader): # Move data to device (may need gradients in future, if so requires_grad=True) invar = Constraint._set_device( invar0, device=self.device, requires_grad=self.requires_grad ) true_outvar = Constraint._set_device( true_outvar0, device=self.device, requires_grad=self.requires_grad ) pred_outvar = self.forward({**invar, **invar_trunk_gpu}) # Collect minibatch info into cpu dictionaries invar_cpu = { key: value + [invar[key].cpu().detach()] for key, value in invar_cpu.items() } true_outvar_cpu = { key: value + [true_outvar[key].cpu().detach()] for key, value in true_outvar_cpu.items() } pred_outvar_cpu = { key: value + [pred_outvar[key].cpu().detach()] for key, value in pred_outvar_cpu.items() } # Concat mini-batch tensors invar_cpu = {key: torch.cat(value) for key, value in invar_cpu.items()} true_outvar_cpu = { key: torch.cat(value) for key, value in true_outvar_cpu.items() } pred_outvar_cpu = { key: torch.cat(value) for key, value in pred_outvar_cpu.items() } # compute losses on cpu losses = DeepONet_Data_Validator._l2_relative_error( true_outvar_cpu, pred_outvar_cpu ) # convert to numpy arrays invar = {k: v.numpy() for k, v in invar_cpu.items()} true_outvar = {k: v.numpy() for k, v in true_outvar_cpu.items()} pred_outvar = {k: v.numpy() for k, v in pred_outvar_cpu.items()} # save batch to vtk file TODO clean this up after graph unroll stuff named_true_outvar = {"true_" + k: v for k, v in true_outvar.items()} named_pred_outvar = {"pred_" + k: v for k, v in pred_outvar.items()} # save batch to vtk/npz file TODO clean this up after graph unroll stuff if "np" in save_filetypes: np.savez( results_dir + name, {**invar, **named_true_outvar, **named_pred_outvar} ) # add tensorboard plots if self.plotter is not None: self.plotter._add_figures( "Validators", name, results_dir, writer, step, self.invar_trunk_plotter, true_outvar, pred_outvar, ) # add tensorboard scalars for k, loss in losses.items(): if TF_SUMMARY: writer.add_scalar("val/" + name + "/" + k, loss, step, new_style=True) else: writer.add_scalar( "Validators/" + name + "/" + k, loss, step, new_style=True ) return losses
© Copyright 2021-2022, NVIDIA. Last updated on Apr 26, 2023.