NVIDIA Modulus Sym v1.0.0
Sym v1.0.0

deeplearning/modulus/modulus-sym-v100/_modules/modulus/sym/hydra/arch.html

Source code for modulus.sym.hydra.arch

# 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.

"""
Architecture/Model configs
"""

from dataclasses import dataclass, field
from hydra.core.config_store import ConfigStore
from omegaconf import MISSING, SI, II
from typing import Any, Union, List, Dict, Tuple


[docs]@dataclass class ModelConf: arch_type: str = MISSING input_keys: Any = MISSING output_keys: Any = MISSING detach_keys: Any = MISSING scaling: Any = None
[docs]@dataclass class AFNOConf(ModelConf): arch_type: str = "afno" img_shape: Tuple[int] = MISSING patch_size: int = 16 embed_dim: int = 256 depth: int = 4 num_blocks: int = 8
[docs]@dataclass class DistributedAFNOConf(ModelConf): arch_type: str = "distributed_afno" img_shape: Tuple[int] = MISSING patch_size: int = 16 embed_dim: int = 256 depth: int = 4 num_blocks: int = 8 channel_parallel_inputs: bool = False channel_parallel_outputs: bool = False
[docs]@dataclass class DeepOConf(ModelConf): arch_type: str = "deeponet" # branch_net: Union[Arch, str], # trunk_net: Union[Arch, str], trunk_dim: Any = None # Union[None, int] branch_dim: Any = None # Union[None, int]
[docs]@dataclass class FNOConf(ModelConf): arch_type: str = "fno" dimension: int = MISSING # decoder_net: Arch nr_fno_layers: int = 4 fno_modes: Any = 16 # Union[int, List[int]] padding: int = 8 padding_type: str = "constant" activation_fn: str = "gelu" coord_features: bool = True
[docs]@dataclass class FourierConf(ModelConf): arch_type: str = "fourier" frequencies: Any = "('axis', [i for i in range(10)])" frequencies_params: Any = "('axis', [i for i in range(10)])" activation_fn: str = "silu" layer_size: int = 512 nr_layers: int = 6 skip_connections: bool = False weight_norm: bool = True adaptive_activations: bool = False
[docs]@dataclass class FullyConnectedConf(ModelConf): arch_type: str = "fully_connected" layer_size: int = 512 nr_layers: int = 6 skip_connections: bool = False activation_fn: str = "silu" adaptive_activations: bool = False weight_norm: bool = True
[docs]@dataclass class ConvFullyConnectedConf(ModelConf): arch_type: str = "conv_fully_connected" layer_size: int = 512 nr_layers: int = 6 skip_connections: bool = False activation_fn: str = "silu" adaptive_activations: bool = False weight_norm: bool = True
[docs]@dataclass class FusedMLPConf(ModelConf): arch_type: str = "fused_fully_connected" layer_size: int = 128 nr_layers: int = 6 activation_fn: str = "sigmoid"
[docs]@dataclass class FusedFourierNetConf(ModelConf): arch_type: str = "fused_fourier" layer_size: int = 128 nr_layers: int = 6 activation_fn: str = "sigmoid" n_frequencies: int = 12
[docs]@dataclass class FusedGridEncodingNetConf(ModelConf): arch_type: str = "fused_hash_encoding" layer_size: int = 128 nr_layers: int = 6 activation_fn: str = "sigmoid" indexing: str = "Hash" n_levels: int = 16 n_features_per_level: int = 2 log2_hashmap_size: int = 19 base_resolution: int = 16 per_level_scale: float = 2.0 interpolation: str = "Smoothstep"
[docs]@dataclass class MultiresolutionHashNetConf(ModelConf): arch_type: str = "hash_encoding" layer_size: int = 64 nr_layers: int = 3 skip_connections: bool = False weight_norm: bool = True adaptive_activations: bool = False bounds: Any = "[(1.0, 1.0), (1.0, 1.0)]" nr_levels: int = 16 nr_features_per_level: int = 2 log2_hashmap_size: int = 19 base_resolution: int = 2 finest_resolution: int = 32
[docs]@dataclass class HighwayFourierConf(ModelConf): arch_type: str = "highway_fourier" frequencies: Any = "('axis', [i for i in range(10)])" frequencies_params: Any = "('axis', [i for i in range(10)])" activation_fn: str = "silu" layer_size: int = 512 nr_layers: int = 6 skip_connections: bool = False weight_norm: bool = True adaptive_activations: bool = False transform_fourier_features: bool = True project_fourier_features: bool = False
[docs]@dataclass class ModifiedFourierConf(ModelConf): arch_type: str = "modified_fourier" frequencies: Any = "('axis', [i for i in range(10)])" frequencies_params: Any = "('axis', [i for i in range(10)])" activation_fn: str = "silu" layer_size: int = 512 nr_layers: int = 6 skip_connections: bool = False weight_norm: bool = True adaptive_activations: bool = False
[docs]@dataclass class MultiplicativeFilterConf(ModelConf): arch_type: str = "multiplicative_fourier" layer_size: int = 512 nr_layers: int = 6 skip_connections: bool = False activation_fn: str = "identity" filter_type: str = "fourier" weight_norm: bool = True input_scale: float = 10.0 gabor_alpha: float = 6.0 gabor_beta: float = 1.0 normalization: Any = ( None # Change to Union[None, Dict[str, Tuple[float, float]]] when supported )
[docs]@dataclass class MultiscaleFourierConf(ModelConf): arch_type: str = "multiscale_fourier" frequencies: Any = field(default_factory=lambda: [32]) frequencies_params: Any = None activation_fn: str = "silu" layer_size: int = 512 nr_layers: int = 6 skip_connections: bool = False weight_norm: bool = True adaptive_activations: bool = False
[docs]@dataclass class Pix2PixConf(ModelConf): arch_type: str = "pix2pix" dimension: int = MISSING conv_layer_size: int = 64 n_downsampling: int = 3 n_blocks: int = 3 scaling_factor: int = 1 batch_norm: bool = True padding_type: str = "reflect" activation_fn: str = "relu"
[docs]@dataclass class SirenConf(ModelConf): arch_type: str = "siren" layer_size: int = 512 nr_layers: int = 6 first_omega: float = 30.0 omega: float = 30.0 normalization: Any = ( None # Change to Union[None, Dict[str, Tuple[float, float]]] when supported )
[docs]@dataclass class SRResConf(ModelConf): arch_type: str = "super_res" large_kernel_size: int = 7 small_kernel_size: int = 3 conv_layer_size: int = 32 n_resid_blocks: int = 8 scaling_factor: int = 8 activation_fn: str = "prelu"
[docs]def register_arch_configs() -> None: # Information regarding multiple config groups # https://hydra.cc/docs/next/patterns/select_multiple_configs_from_config_group/ cs = ConfigStore.instance() cs.store( group="arch", name="fused_fully_connected", node={"fused_fully_connected": FusedMLPConf()}, ) cs.store( group="arch", name="fused_fourier", node={"fused_fourier": FusedFourierNetConf()}, ) cs.store( group="arch", name="fused_hash_encoding", node={"fused_hash_encoding": FusedGridEncodingNetConf()}, ) cs.store( group="arch", name="fully_connected", node={"fully_connected": FullyConnectedConf()}, ) cs.store( group="arch", name="conv_fully_connected", node={"conv_fully_connected": ConvFullyConnectedConf()}, ) cs.store( group="arch", name="fourier", node={"fourier": FourierConf()}, ) cs.store( group="arch", name="highway_fourier", node={"highway_fourier": HighwayFourierConf()}, ) cs.store( group="arch", name="modified_fourier", node={"modified_fourier": ModifiedFourierConf()}, ) cs.store( group="arch", name="multiplicative_fourier", node={"multiplicative_fourier": MultiplicativeFilterConf()}, ) cs.store( group="arch", name="multiscale_fourier", node={"multiscale_fourier": MultiscaleFourierConf()}, ) cs.store( group="arch", name="siren", node={"siren": SirenConf()}, ) cs.store( group="arch", name="hash_encoding", node={"hash_encoding": MultiresolutionHashNetConf()}, ) cs.store( group="arch", name="fno", node={"fno": FNOConf()}, ) cs.store( group="arch", name="afno", node={"afno": AFNOConf()}, ) cs.store( group="arch", name="distributed_afno", node={"distributed_afno": DistributedAFNOConf()}, ) cs.store( group="arch", name="deeponet", node={"deeponet": DeepOConf()}, ) cs.store( group="arch", name="super_res", node={"super_res": SRResConf()}, ) cs.store( group="arch", name="pix2pix", node={"pix2pix": Pix2PixConf()}, ) # Schemas for extending models # Info: https://hydra.cc/docs/next/patterns/extending_configs/ cs.store( group="arch", name="fully_connected_cfg", node=FullyConnectedConf, ) cs.store( group="arch", name="conv_fully_connected_cfg", node=ConvFullyConnectedConf, ) cs.store( group="arch", name="fused_mlp_cfg", node=FusedMLPConf, ) cs.store( group="arch", name="fused_fourier_net_cfg", node=FusedFourierNetConf, ) cs.store( group="arch", name="fused_grid_encoding_net_cfg", node=FusedGridEncodingNetConf, ) cs.store( group="arch", name="fourier_cfg", node=FourierConf, ) cs.store( group="arch", name="highway_fourier_cfg", node=HighwayFourierConf, ) cs.store( group="arch", name="modified_fourier_cfg", node=ModifiedFourierConf, ) cs.store( group="arch", name="multiplicative_fourier_cfg", node=MultiplicativeFilterConf, ) cs.store( group="arch", name="multiscale_fourier_cfg", node=MultiscaleFourierConf, ) cs.store( group="arch", name="siren_cfg", node=SirenConf, ) cs.store( group="arch", name="hash_net_cfg", node=MultiresolutionHashNetConf, ) cs.store( group="arch", name="fno_cfg", node=FNOConf, ) cs.store( group="arch", name="afno_cfg", node=AFNOConf, ) cs.store( group="arch", name="distributed_afno_cfg", node=DistributedAFNOConf, ) cs.store( group="arch", name="deeponet_cfg", node=DeepOConf, ) cs.store( group="arch", name="super_res_cfg", node=SRResConf, ) cs.store( group="arch", name="pix2pix_cfg", node=Pix2PixConf, )
© Copyright 2023, NVIDIA Modulus Team. Last updated on Aug 8, 2023.