deeplearning/modulus/modulus-sym-v120/_modules/modulus/sym/models/multiplicative_filter_net.html
Source code for modulus.sym.models.multiplicative_filter_net
# 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 enum
from typing import Optional, List, Dict, Tuple, Union
import torch
import torch.nn as nn
from torch import Tensor
import modulus.sym.models.layers as layers
from modulus.sym.models.arch import Arch
from modulus.sym.models.layers import Activation
from modulus.sym.key import Key
from modulus.sym.constants import NO_OP_NORM
class FilterTypeMeta(enum.EnumMeta):
def __getitem__(self, name):
try:
return super().__getitem__(name.upper())
except (KeyError) as error:
raise KeyError(f"Invalid activation function {name}")
class FilterType(enum.Enum, metaclass=FilterTypeMeta):
FOURIER = enum.auto()
GABOR = enum.auto()
[docs]class MultiplicativeFilterNetArch(Arch):
"""
Multiplicative Filter Net with Activations
Reference: Fathony, R., Sahu, A.K., AI, A.A., Willmott, D. and Kolter, J.Z., MULTIPLICATIVE FILTER NETWORKS.
Parameters
----------
input_keys : List[Key]
Input key list
output_keys : List[Key]
Output key list
detach_keys : List[Key], optional
List of keys to detach gradients, by default []
layer_size : int = 512
Layer size for every hidden layer of the model.
nr_layers : int = 6
Number of hidden layers of the model.
skip_connections : bool = False
If true then apply skip connections every 2 hidden layers.
activation_fn : layers.Activation = layers.Activation.SILU
Activation function used by network.
filter_type : FilterType = FilterType.FOURIER
Filter type for multiplicative filter network, (Fourier or Gabor).
weight_norm : bool = True
Use weight norm on fully connected layers.
input_scale : float = 10.0
Scale inputs for multiplicative filters.
gabor_alpha : float = 6.0
Alpha value for Gabor filter.
gabor_beta : float = 1.0
Beta value for Gabor filter.
normalization : Optional[Dict[str, Tuple[float, float]]] = None
Normalization of input to network.
"""
def __init__(
self,
input_keys: List[Key],
output_keys: List[Key],
detach_keys: List[Key] = [],
layer_size: int = 512,
nr_layers: int = 6,
skip_connections: bool = False,
activation_fn=layers.Activation.IDENTITY,
filter_type: Union[FilterType, str] = FilterType.FOURIER,
weight_norm: bool = True,
input_scale: float = 10.0,
gabor_alpha: float = 6.0,
gabor_beta: float = 1.0,
normalization: Optional[Dict[str, Tuple[float, float]]] = None,
) -> None:
super().__init__(
input_keys=input_keys, output_keys=output_keys, detach_keys=detach_keys
)
in_features = sum(self.input_key_dict.values())
out_features = sum(self.output_key_dict.values())
self.nr_layers = nr_layers
self.skip_connections = skip_connections
if isinstance(filter_type, str):
filter_type = FilterType[filter_type]
if filter_type == FilterType.FOURIER:
self.first_filter = layers.FourierFilter(
in_features=in_features,
layer_size=layer_size,
nr_layers=nr_layers,
input_scale=input_scale,
)
elif filter_type == FilterType.GABOR:
self.first_filter = layers.GaborFilter(
in_features=in_features,
layer_size=layer_size,
nr_layers=nr_layers,
input_scale=input_scale,
alpha=gabor_alpha,
beta=gabor_beta,
)
else:
raise ValueError
self.filters = nn.ModuleList()
self.fc_layers = nn.ModuleList()
for i in range(nr_layers):
self.fc_layers.append(
layers.FCLayer(
in_features=layer_size,
out_features=layer_size,
activation_fn=activation_fn,
weight_norm=weight_norm,
)
)
if filter_type == FilterType.FOURIER:
self.filters.append(
layers.FourierFilter(
in_features=in_features,
layer_size=layer_size,
nr_layers=nr_layers,
input_scale=input_scale,
)
)
elif filter_type == FilterType.GABOR:
self.filters.append(
layers.GaborFilter(
in_features=in_features,
layer_size=layer_size,
nr_layers=nr_layers,
input_scale=input_scale,
alpha=gabor_alpha,
beta=gabor_beta,
)
)
else:
raise ValueError
self.final_layer = layers.FCLayer(
in_features=layer_size,
out_features=out_features,
activation_fn=layers.Activation.IDENTITY,
weight_norm=False,
activation_par=None,
)
self.normalization: Optional[Dict[str, Tuple[float, float]]] = normalization
# iterate input keys and add NO_OP_NORM if it is not specified
if self.normalization is not None:
for key in self.input_key_dict:
if key not in self.normalization:
self.normalization[key] = NO_OP_NORM
self.register_buffer(
"normalization_tensor",
self._get_normalization_tensor(self.input_key_dict, self.normalization),
persistent=False,
)
def _tensor_forward(self, x: Tensor) -> Tensor:
x = self._tensor_normalize(x, self.normalization_tensor)
x = self.process_input(
x, self.input_scales_tensor, input_dict=self.input_key_dict, dim=-1
)
res = self.first_filter(x)
res_skip: Optional[Tensor] = None
for i, (fc_layer, filter) in enumerate(zip(self.fc_layers, self.filters)):
res_fc = fc_layer(res)
res_filter = filter(x)
res = res_fc * res_filter
if self.skip_connections and i % 2 == 0:
if res_skip is not None:
res, res_skip = res + res_skip, res
else:
res_skip = res
x = self.final_layer(res)
x = self.process_output(x, self.output_scales_tensor)
return x
[docs] def forward(self, in_vars: Dict[str, Tensor]) -> Dict[str, Tensor]:
x = self.concat_input(
in_vars,
self.input_key_dict.keys(),
detach_dict=self.detach_key_dict,
dim=-1,
)
y = self._tensor_forward(x)
return self.split_output(y, self.output_key_dict, dim=-1)def _dict_forward(self, in_vars: Dict[str, Tensor]) -> Dict[str, Tensor]:
"""
This is the original forward function, left here for the correctness test.
"""
x = self.prepare_input(
self._normalize(in_vars, self.normalization),
self.input_key_dict.keys(),
detach_dict=self.detach_key_dict,
dim=-1,
input_scales=self.input_scales,
)
res = self.first_filter(x)
res_skip: Optional[Tensor] = None
for i, (fc_layer, filter) in enumerate(zip(self.fc_layers, self.filters)):
res_fc = fc_layer(res)
res_filter = filter(x)
res = res_fc * res_filter
if self.skip_connections and i % 2 == 0:
if res_skip is not None:
res, res_skip = res + res_skip, res
else:
res_skip = res
res = self.final_layer(res)
return self.prepare_output(
res, self.output_key_dict, dim=-1, output_scales=self.output_scales
)
def _normalize(
self,
in_vars: Dict[str, Tensor],
norms: Optional[Dict[str, Tuple[float, float]]],
) -> Dict[str, Tensor]:
if norms is None:
return in_vars
normalized_in_vars = {}
for k, v in in_vars.items():
if k in norms:
v = (v - norms[k][0]) / (norms[k][1] - norms[k][0])
v = 2 * v - 1
normalized_in_vars[k] = v
return normalized_in_vars