Convenient helper nn.Module to create a stack of linear layers. Args: num_layers: Total number of linear layers to create. dim_input: Feature dimension of the input tensor, which is :math:`(*, C_{in})`. dim_output:
(
self,
num_layers: int,
dim_input: int,
dim_output: int,
dim_features: T.Union[T.Sequence[int], int],
nonlinearity: str = "leaky_relu",
add_norm_layer: bool = False,
norm_fun: T.Callable = nn.LayerNorm,
dropout_prob: float = 0.0,
output_add_nonlinearity: bool = False,
)
| 108 | """ |
| 109 | |
| 110 | def __init__( |
| 111 | self, |
| 112 | num_layers: int, |
| 113 | dim_input: int, |
| 114 | dim_output: int, |
| 115 | dim_features: T.Union[T.Sequence[int], int], |
| 116 | nonlinearity: str = "leaky_relu", |
| 117 | add_norm_layer: bool = False, |
| 118 | norm_fun: T.Callable = nn.LayerNorm, |
| 119 | dropout_prob: float = 0.0, |
| 120 | output_add_nonlinearity: bool = False, |
| 121 | ): |
| 122 | """ |
| 123 | Convenient helper nn.Module to create a stack of linear layers. |
| 124 | |
| 125 | Args: |
| 126 | num_layers: |
| 127 | Total number of linear layers to create. |
| 128 | dim_input: |
| 129 | Feature dimension of the input tensor, which is :math:`(*, C_{in})`. |
| 130 | dim_output: |
| 131 | Feature dimension of the output tensor, which is :math:`(*, C_{out})`. |
| 132 | dim_features: |
| 133 | An integer if all layers share the same feature dimension, |
| 134 | or a list of num_layer-1 integers, one for each layer except the last layer. |
| 135 | nonlinearity: |
| 136 | Nonlinearity used after each linear layer (except the last layer if output_add_nonlinearity is False). |
| 137 | Choose from: `leaky_relu`, `relu`, `tanh`, `sigmoid`, `silu`, `swish` |
| 138 | Note that silu (swish) is supported in pytorch version >= 1.7.0. |
| 139 | add_norm_layer: |
| 140 | Whether to add normalization layers between linear layers |
| 141 | norm_fun: |
| 142 | Callable function used to normalize the output of linear layer (before nonlinearity). |
| 143 | It should be a function that takes dim_feature as input. |
| 144 | For example, you can pass `nn.LayerNorm`. |
| 145 | If you want to control additional functionality like the eps and elementwise_affine of nn.LayerNorm, |
| 146 | you can pass a lambda function: |
| 147 | lambda dim: torch.nn.LayerNorm(dim, eps=1e-5, elementwise_affine=False) |
| 148 | dropout_prob: |
| 149 | Dropout probability added after nonlinearity. If 0, no dropout layer is added. |
| 150 | output_add_nonlinearity: |
| 151 | Whether to add nonlinearity (norm_layer, and dropout) at the last layer |
| 152 | |
| 153 | Note that the order of the layers is: |
| 154 | Linear -> normalization (if add_norm_layer) -> nonlinearity -> dropout. |
| 155 | |
| 156 | """ |
| 157 | super().__init__() |
| 158 | self.dim_input = dim_input |
| 159 | self.num_layers = num_layers |
| 160 | self.dim_output = dim_output |
| 161 | self.add_norm_layer = add_norm_layer |
| 162 | self.linear_bias = not self.add_norm_layer # if added normalization layer, no need to learn bias |
| 163 | self.norm_fun = norm_fun |
| 164 | self.dropout_prob = dropout_prob |
| 165 | self.output_add_nonlinearity = output_add_nonlinearity |
| 166 | |
| 167 | # check dim_features |
nothing calls this directly
no test coverage detected