# Copyright (c) 2023 Predibase, Inc., 2019 Uber Technologies, Inc. # # 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 logging from copy import deepcopy import torch from torch.nn import Dropout, Linear, ModuleList from ludwig.modules.normalization_modules import create_norm_layer from ludwig.utils.torch_utils import activations, initializer_registry, LudwigModule logger = logging.getLogger(__name__) class FCLayer(LudwigModule): """A torch.nn.Linear wrapper that declares input and output shapes, and enables the customization of: 1. how weights and biases are initialized 2. normalization (layer and batch) 3. activations 4. dropout """ def __init__( self, input_size: int, input_rank: int = 2, output_size: int = 256, use_bias: bool = True, weights_initializer: str = "xavier_uniform", bias_initializer: str = "zeros", norm: str | None = None, norm_params: dict | None = None, activation: str = "relu", dropout: float = 0, ): super().__init__() self.layers = ModuleList() self.input_size = input_size self.output_size = output_size fc = Linear(in_features=input_size, out_features=output_size, bias=use_bias) self.layers.append(fc) weights_initializer = initializer_registry[weights_initializer] weights_initializer(fc.weight) if use_bias: bias_initializer = initializer_registry[bias_initializer] bias_initializer(fc.bias) if norm is not None: norm_params = norm_params or {} self.layers.append(create_norm_layer(norm, input_rank, output_size, **norm_params)) # Dict for activation objects in pytorch? self.layers.append(activations[activation]()) if dropout > 0: self.layers.append(Dropout(dropout)) def forward(self, inputs, mask=None): hidden = inputs for layer in self.layers: hidden = layer(hidden) return hidden @property def input_shape(self) -> torch.Size: return torch.Size([self.input_size]) @property def output_shape(self) -> torch.Size: return torch.Size([self.output_size]) class FCStack(LudwigModule): """A stack of FCLayers. The specification of each FCLayer is specified by the `layers` dictionary parameter, whose keys correspond with an FCLayer's constructor arguments, i.e. [ {"input_size": 2, "output_size": 4}, {"output_size": 4, "use_bias": False}, ] `default_*` parameters dictate default values to use for each FCLayer, if not specified by `layers`. If `layers` is `None`, then a stack of size `num_layers` of `FCLayer`s configured with all of the `default_*` parameters is used. If `layers` is None and `num_layers` is 0, then there are no fully connected layers and this module serves as a trivial passthrough. """ def __init__( self, first_layer_input_size: int, layers: list[dict] | None = None, num_layers: int = 1, default_input_rank: int = 2, default_output_size: int = 256, default_use_bias: bool = True, default_weights_initializer: str = "xavier_uniform", default_bias_initializer: str = "zeros", default_norm: str | None = None, default_norm_params: dict | None = None, default_activation: str = "relu", default_dropout: float = 0, residual: bool = False, **kwargs, ): super().__init__() self.input_size = first_layer_input_size self.norm_layer = None if default_norm is not None: norm_params = default_norm_params or {} self.norm_layer = create_norm_layer(default_norm, default_input_rank, self.input_size, **norm_params) self.dropout = None if default_dropout > 0: self.dropout = torch.nn.Dropout(default_dropout) if layers is None: self.layers = [] for i in range(num_layers): self.layers.append({}) else: # deep copy the layer definitions so that we don't modify the original self.layers = deepcopy(layers) if len(self.layers) > 0 and "input_size" not in self.layers[0]: self.layers[0]["input_size"] = first_layer_input_size for i, layer in enumerate(self.layers): if i != 0: layer["input_size"] = self.layers[i - 1]["output_size"] if "input_rank" not in layer: layer["input_rank"] = default_input_rank if "output_size" not in layer: layer["output_size"] = default_output_size if "use_bias" not in layer: layer["use_bias"] = default_use_bias if "weights_initializer" not in layer: layer["weights_initializer"] = default_weights_initializer if "bias_initializer" not in layer: layer["bias_initializer"] = default_bias_initializer if "norm" not in layer: layer["norm"] = default_norm if "norm_params" not in layer: layer["norm_params"] = default_norm_params if "activation" not in layer: layer["activation"] = default_activation if "dropout" not in layer: layer["dropout"] = default_dropout self.stack = ModuleList() for i, layer in enumerate(self.layers): self.stack.append( FCLayer( input_size=layer["input_size"], input_rank=layer["input_rank"], output_size=layer["output_size"], use_bias=layer["use_bias"], weights_initializer=layer["weights_initializer"], bias_initializer=layer["bias_initializer"], norm=layer["norm"], norm_params=layer["norm_params"], activation=layer["activation"], dropout=layer["dropout"], ) ) self.residual = residual def forward(self, inputs, mask=None): hidden = inputs if self.norm_layer is not None: hidden = self.norm_layer(hidden) if self.dropout is not None: hidden = self.dropout(hidden) prev_fc_layer_size = self.input_size for layer in self.stack: out = layer(hidden) if self.residual and layer.output_size == prev_fc_layer_size: hidden = hidden + out else: hidden = out prev_fc_layer_size = layer.layers[0].out_features return hidden @property def num_layers(self) -> int: return len(self.layers) @property def input_shape(self) -> torch.Size: return torch.Size([self.input_size]) @property def output_shape(self) -> torch.Size: if len(self.stack) > 0: return self.stack[-1].output_shape return torch.Size([self.input_size])