torchcnnbuilder.builder

  1import math
  2from collections import OrderedDict
  3from typing import List, Optional, Sequence, Tuple, Union
  4
  5import torch.nn as nn
  6
  7from torchcnnbuilder._constants import (
  8    DEFAULT_CONV_PARAMS,
  9    DEFAULT_TRANSPOSE_CONV_PARAMS,
 10)
 11from torchcnnbuilder._formulas import _select_conv_calc
 12from torchcnnbuilder._utils import (
 13    _double_params,
 14    _select_adaptive_pooling_dimension,
 15    _select_conv_dimension,
 16    _select_norm_dimension,
 17    _set_conv_params,
 18    _triple_params,
 19)
 20from torchcnnbuilder._validation import (
 21    _validate_available_layers,
 22    _validate_build_transpose_convolve_init,
 23    _validate_channel_growth_rate_param,
 24    _validate_difference_in_dimensions,
 25    _validate_input_size_is_not_none,
 26    _validate_max_channels_number,
 27    _validate_min_channels_number,
 28    _validate_range_step,
 29)
 30from torchcnnbuilder.latent import LatentSpaceModule
 31
 32
 33# ------------------------------------
 34# CNN Builder class
 35# ------------------------------------
 36class Builder:
 37    """
 38    A class for creating СNN architectures
 39    """
 40
 41    def __init__(
 42        self,
 43        input_size: Optional[Sequence[int]] = None,
 44        minimum_feature_map_size: Union[Sequence[int], int] = 5,
 45        max_channels: int = 512,
 46        min_channels: int = 1,
 47        activation_function: nn.Module = nn.ReLU(inplace=True),
 48        finish_activation_function: Union[Optional[nn.Module], str] = None,
 49    ) -> None:
 50        """
 51        Initializes the Builder instance.
 52
 53        Args:
 54            input_size (Optional[Sequence[int]], optional):
 55                Input size of the input tensor. Necessary for creating
 56                convolution sequences. Defaults to None.
 57            minimum_feature_map_size (Union[Sequence[int], int], optional):
 58                Minimum feature map size. Defaults to 5.
 59            max_channels (int, optional):
 60                Maximum number of layers after any convolution. Defaults to 512.
 61            min_channels (int, optional):
 62                Minimum number of layers after any convolution. Defaults to 1.
 63            activation_function (nn.Module, optional):
 64                Activation function. Defaults to nn.ReLU(inplace=True).
 65            finish_activation_function (Union[Optional[nn.Module], str], optional):
 66                Last activation function, can be the same as activation_function
 67                (use string 'same' for that). Defaults to None.
 68
 69        Raises:
 70            ValueError: If input_size is not a valid shape.
 71        """
 72
 73        if input_size is None:
 74            self.input_size = input_size
 75        else:
 76            self.input_size = tuple(i for i in input_size)
 77
 78            if len(self.input_size) == 1:
 79                self.minimum_feature_map_size = (
 80                    (minimum_feature_map_size,)
 81                    if isinstance(minimum_feature_map_size, int)
 82                    else minimum_feature_map_size
 83                )
 84            if len(self.input_size) == 2:
 85                self.minimum_feature_map_size = (
 86                    _double_params(
 87                        minimum_feature_map_size,
 88                    )
 89                    if isinstance(minimum_feature_map_size, int)
 90                    else minimum_feature_map_size
 91                )
 92            if len(self.input_size) == 3:
 93                self.minimum_feature_map_size = (
 94                    _triple_params(
 95                        minimum_feature_map_size,
 96                    )
 97                    if isinstance(minimum_feature_map_size, int)
 98                    else minimum_feature_map_size
 99                )
100
101        self.max_channels = max_channels
102        self._initial_max_channels = max_channels
103
104        self.min_channels = min_channels
105        self._initial_min_channels = min_channels
106
107        self._default_convolve_params = DEFAULT_CONV_PARAMS
108        self._default_transpose_params = DEFAULT_TRANSPOSE_CONV_PARAMS
109
110        # finish_activation_function can be str 'same' which equals to activation_function
111        self.activation_function = activation_function
112        self.finish_activation_function = finish_activation_function
113
114        self._conv_channels = None
115        self._transpose_conv_channels = None
116
117        self._conv_layers = None
118        self._transpose_conv_layers = None
119
120    @property
121    def conv_channels(self) -> Optional[List[int]]:
122        """Gets the convolutional channels.
123
124        Returns:
125            A list of convolutional channel sizes or None
126            if not initialized.
127        """
128        return self._conv_channels
129
130    @property
131    def transpose_conv_channels(self) -> Optional[List[int]]:
132        """Gets the transposed convolutional channels.
133
134        Returns:
135            A list of transposed convolutional channel sizes
136            or None if not initialized.
137        """
138        return self._transpose_conv_channels
139
140    @property
141    def conv_layers(self) -> Optional[List[Tuple[int, ...]]]:
142        """Gets the convolutional layers.
143
144        Returns:
145            A list of tuples representing convolutional layer configurations or None if not initialized.
146        """
147        return self._conv_layers
148
149    @property
150    def transpose_conv_layers(self) -> Optional[List[Tuple[int, ...]]]:
151        """Gets the transposed convolutional layers.
152
153        Returns:
154            A list of tuples representing transposed convolutional layer configurations or None if not
155            initialized.
156        """
157        return self._transpose_conv_layers
158
159    def build_convolve_block(
160        self,
161        in_channels: int,
162        out_channels: int,
163        params: Optional[dict] = None,
164        normalization: Optional[str] = None,
165        sub_blocks: int = 1,
166        p: float = 0.5,
167        inplace: bool = False,
168        eps: float = 1e-5,
169        momentum: Optional[float] = 0.1,
170        affine: bool = True,
171        conv_dim: int = 2,
172    ) -> nn.Sequential:
173        """Builds a single block of convolution layers.
174
175        This method creates a sequential block that consists of convolutional layers,
176        optional normalization layers, and an activation function.
177
178        Args:
179            in_channels (int): Number of channels in the input image.
180            out_channels (int): Number of channels produced by the convolution.
181            params (Optional[dict], optional): Convolutional layer parameters
182                (for nn.ConvNd). Defaults to None.
183            normalization (Optional[str], optional): Type of normalization to apply.
184                Options are 'dropout', 'instancenorm', and 'batchnorm'. Defaults to None.
185            sub_blocks (int, optional): Number of convolutional layers within the block.
186                Defaults to 1.
187            p (float, optional): Probability of an element being zeroed for dropout/instancenorm.
188                Defaults to 0.5.
189            inplace (bool, optional): If True, performs the operation in-place
190                for dropout/instancenorm. Defaults to False.
191            eps (float, optional): A value added to the denominator for numerical stability
192                (used in batchnorm/instancenorm). Defaults to 1e-5.
193            momentum (Optional[float], optional): Momentum for running_mean or running_var
194                computation (used in batchnorm). If None, a cumulative moving average is used.
195                Defaults to 0.1.
196            affine (bool, optional): If True, the module has learnable affine parameters
197                (used in batchnorm). Defaults to True.
198            conv_dim (int, optional): The dimension of the convolutional operation (2 for
199                2D convolution, 3 for 3D convolution). Defaults to 2.
200
201        Returns:
202            A sequential block containing convolutional layers,
203            optional normalization layers, and an activation function.
204        """
205        params = _set_conv_params(default_params=self._default_convolve_params, params=params)
206        convolution = _select_conv_dimension(conv_dim=conv_dim)
207
208        if sub_blocks > 1:
209            kernel_size = params["kernel_size"]
210            kernel_size = kernel_size if isinstance(kernel_size, int) else kernel_size[0]
211            params["padding"] = kernel_size // 2
212            params["stride"] = 1
213
214        blocks = []
215        for i in range(sub_blocks):
216            block = []
217
218            conv = convolution(in_channels=in_channels, out_channels=out_channels, **params)
219            in_channels = out_channels
220            block.append(conv)
221
222            if normalization:
223                norm = _select_norm_dimension(conv_dim=conv_dim, normalization=normalization)
224
225                if normalization in ("batchnorm", "instancenorm"):
226                    norm = norm(
227                        num_features=out_channels,
228                        eps=eps,
229                        momentum=momentum,
230                        affine=affine,
231                    )
232
233                if normalization == "dropout":
234                    norm = norm(p=p, inplace=inplace)
235
236                block.append(norm)
237
238            activation_function = self.activation_function
239            block.append(activation_function)
240
241            if sub_blocks > 1:
242                block = nn.Sequential(*block)
243                blocks.append((f"sub-block {i + 1}", block))
244            else:
245                blocks.extend(block)
246
247        if sub_blocks > 1:
248            return nn.Sequential(OrderedDict(blocks))
249
250        return nn.Sequential(*blocks)
251
252    def build_convolve_sequence(
253        self,
254        n_layers: int,
255        in_channels: int = 1,
256        params: Optional[dict] = None,
257        normalization: Optional[str] = None,
258        sub_blocks: int = 1,
259        p: float = 0.5,
260        inplace: bool = False,
261        eps: float = 1e-5,
262        momentum: Optional[float] = 0.1,
263        affine: bool = True,
264        ratio: float = 2.0,
265        start: int = 32,
266        channel_growth_rate: str = "exponential",
267        conv_dim: int = 2,
268    ) -> nn.Sequential:
269        """Builds a sequence of convolution blocks.
270
271        This method constructs a sequential block of convolutional layers,
272        with optional normalization and multiple sub-blocks per layer.
273
274        Args:
275            n_layers (int): Number of convolution layers in the encoder part.
276            in_channels (int, optional): Number of channels in the first input tensor.
277                Defaults to 1.
278            params (Optional[dict], optional): Convolutional layer parameters
279                (for nn.ConvNd). Defaults to None.
280            normalization (Optional[str], optional): Type of normalization to apply.
281                Options are 'dropout', 'instancenorm', and 'batchnorm'. Defaults to None.
282            sub_blocks (int, optional): Number of convolutions within each layer.
283                Defaults to 1.
284            p (float, optional): Probability of an element being zeroed for dropout.
285                Defaults to 0.5.
286            inplace (bool, optional): If True, performs the operation in-place
287                for dropout. Defaults to False.
288            eps (float, optional): A value added to the denominator for numerical stability
289                (used in batchnorm/instancenorm). Defaults to 1e-5.
290            momentum (Optional[float], optional): Momentum for running_mean or running_var
291                computation (used in batchnorm). If None, a cumulative moving average is used.
292                Defaults to 0.1.
293            affine (bool, optional): If True, the module has learnable affine parameters
294                (used in batchnorm). Defaults to True.
295            ratio (float, optional): Multiplier for the geometric progression of increasing
296                channels (feature maps). Used for 'channel_growth_rate' as 'exponential'
297                or 'power'. Defaults to 2.0.
298            start (int, optional): Starting position of the geometric progression
299                when 'channel_growth_rate' is set to 'exponential'. Defaults to 32.
300            channel_growth_rate (str, optional): Method for calculating the number of
301                feature maps. Options include 'exponential', 'proportion', 'linear',
302                'power', and 'constant'. Defaults to 'exponential'.
303            conv_dim (int, optional): The dimension of the convolutional operation.
304                Defaults to 2.
305
306        Returns:
307            A sequential block containing the specified number of
308            convolutional layers.
309        """
310        _validate_input_size_is_not_none(self.input_size)
311        params = _set_conv_params(default_params=self._default_convolve_params, params=params)
312        conv_out = _select_conv_calc(conv_dim=conv_dim)
313
314        modules = []
315        input_layer_size_list = [self.input_size]
316        input_channels_count_list = self._calc_out_channels(
317            in_size=self.input_size,
318            in_channels=in_channels,
319            n_layers=n_layers,
320            ratio=ratio,
321            start=start,
322            channel_growth_rate=channel_growth_rate,
323        )
324
325        for layer in range(n_layers):
326            input_layer_size = input_layer_size_list[-1]
327
328            _validate_difference_in_dimensions(self.input_size, conv_dim)
329            _validate_available_layers(layer, input_layer_size, self.minimum_feature_map_size)
330            _validate_max_channels_number(layer, input_channels_count_list, self.max_channels)
331            _validate_min_channels_number(layer, input_channels_count_list, self.min_channels)
332
333            in_channels = input_channels_count_list[layer]
334            out_channels = input_channels_count_list[layer + 1]
335
336            out_layer_size = conv_out(input_size=input_layer_size, **params)
337            input_layer_size_list.append(out_layer_size)
338
339            convolve_block = self.build_convolve_block(
340                in_channels=in_channels,
341                out_channels=out_channels,
342                normalization=normalization,
343                sub_blocks=sub_blocks,
344                p=p,
345                inplace=inplace,
346                eps=eps,
347                momentum=momentum,
348                affine=affine,
349                params=params,
350                conv_dim=conv_dim,
351            )
352
353            modules.append((f"conv {layer + 1}", convolve_block))
354
355        self._conv_channels = input_channels_count_list
356        self._conv_layers = input_layer_size_list
357        return nn.Sequential(OrderedDict(modules))
358
359    def build_transpose_convolve_block(
360        self,
361        in_channels: int,
362        out_channels: int,
363        params: Optional[dict] = None,
364        normalization: Optional[str] = None,
365        sub_blocks: int = 1,
366        p: float = 0.5,
367        inplace: bool = False,
368        eps: float = 1e-5,
369        momentum: Optional[float] = 0.1,
370        affine: bool = True,
371        last_block: bool = False,
372        conv_dim: int = 2,
373    ) -> nn.Sequential:
374        """Builds a single block of transposed convolution layers.
375
376        This method constructs a sequential block of transposed convolutional layers,
377        with optional normalization and multiple sub-blocks per layer.
378
379        Args:
380            in_channels (int): Number of channels in the input image.
381            out_channels (int): Number of channels produced by the transposed convolution.
382            params (Optional[dict], optional): Parameters for the transposed convolutional layer
383                (for nn.ConvTranspose2d). Defaults to None.
384            normalization (Optional[str], optional): Type of normalization to apply.
385                Options are 'dropout', 'instancenorm', and 'batchnorm'. Defaults to None.
386            sub_blocks (int, optional): Number of convolutions within each layer.
387                Defaults to 1.
388            p (float, optional): Probability of an element being zeroed for dropout.
389                Defaults to 0.5.
390            inplace (bool, optional): If True, performs the operation in-place
391                for dropout. Defaults to False.
392            eps (float, optional): A value added to the denominator for numerical stability
393                (used in batchnorm/instancenorm). Defaults to 1e-5.
394            momentum (Optional[float], optional): Momentum for running_mean or running_var
395                computation (used in batchnorm). If None, a cumulative moving average is used.
396                Defaults to 0.1.
397            affine (bool, optional): If True, the module has learnable affine parameters
398                (used in batchnorm). Defaults to True.
399            last_block (bool, optional): If True, no activation function is applied after
400                the transposed convolution. Defaults to False.
401            conv_dim (int, optional): The dimension of the convolutional operation.
402                Defaults to 2.
403
404        Returns:
405            A sequential block containing the specified transposed
406            convolutional layers, possibly including normalization and activation functions.
407        """
408        params = _set_conv_params(default_params=self._default_transpose_params, params=params)
409        convolution = _select_conv_dimension(conv_dim=conv_dim, transpose=True)
410
411        if sub_blocks > 1:
412            kernel_size = params["kernel_size"]
413            kernel_size = kernel_size if isinstance(kernel_size, int) else kernel_size[0]
414            params["padding"] = kernel_size // 2
415            params["stride"] = 1
416
417        blocks = []
418        last_out_channels = out_channels
419        for i in range(sub_blocks):
420            block = []
421
422            out_channels = last_out_channels if i == sub_blocks - 1 else in_channels
423            conv = convolution(in_channels=in_channels, out_channels=out_channels, **params)
424            block.append(conv)
425
426            if normalization:
427                norm = _select_norm_dimension(conv_dim=conv_dim, normalization=normalization)
428
429                if normalization in ("batchnorm", "instancenorm"):
430                    norm = norm(
431                        num_features=out_channels,
432                        eps=eps,
433                        momentum=momentum,
434                        affine=affine,
435                    )
436
437                if normalization == "dropout":
438                    norm = norm(p=p, inplace=inplace)
439
440                block.append(norm)
441
442            activation_function = self.activation_function
443            if last_block and i == sub_blocks - 1:
444                if self.finish_activation_function == "same":
445                    block.append(activation_function)
446                elif self.finish_activation_function:
447                    block.append(self.finish_activation_function)
448            else:
449                block.append(activation_function)
450
451            if sub_blocks > 1:
452                block = nn.Sequential(*block)
453                blocks.append((f"transpose sub-block {i + 1}", block))
454            else:
455                blocks.extend(block)
456
457        if sub_blocks > 1:
458            return nn.Sequential(OrderedDict(blocks))
459
460        return nn.Sequential(*blocks)
461
462    def build_transpose_convolve_sequence(
463        self,
464        n_layers: int,
465        in_channels: Optional[int] = None,
466        out_channels: int = 1,
467        out_size: Optional[tuple] = None,
468        params: Optional[dict] = None,
469        normalization: Optional[str] = None,
470        sub_blocks: int = 1,
471        p: float = 0.5,
472        inplace: bool = False,
473        eps: float = 1e-5,
474        momentum: Optional[float] = 0.1,
475        affine: bool = True,
476        ratio: float = 2.0,
477        channel_growth_rate: str = "exponential",
478        conv_dim: int = 2,
479        adaptive_pool: str = "avgpool",
480    ) -> nn.Sequential:
481        """Builds a sequence of transposed convolution blocks.
482
483        This method constructs a sequential layer of transposed convolution blocks,
484        allowing for customization through normalization, sub-blocks, and other parameters.
485
486        Args:
487            n_layers (int): Number of transposed convolution layers to create.
488            in_channels (Optional[int], optional): Number of channels in the first input tensor.
489                Defaults to None, which will use the last value from _conv_channels if available.
490            out_channels (int, optional): Number of channels after the transposed convolution sequence.
491                Defaults to 1.
492            out_size (Optional[tuple], optional): Desired output size after the transposed convolution sequence.
493                Defaults to None, which uses the input size.
494            params (Optional[dict], optional): Parameters for the transposed convolutional layer
495                (for nn.ConvTranspose2d). Defaults to None.
496            normalization (Optional[str], optional): Type of normalization to apply.
497                Options include 'dropout', 'instancenorm', and 'batchnorm'. Defaults to None.
498            sub_blocks (int, optional): Number of transposed convolutions within each layer.
499                Defaults to 1.
500            p (float, optional): Probability of an element being zeroed for dropout.
501                Defaults to 0.5.
502            inplace (bool, optional): If True, performs the operation in-place
503                for dropout. Defaults to False.
504            eps (float, optional): A value added to the denominator for numerical stability
505                (used in batchnorm/instancenorm). Defaults to 1e-5.
506            momentum (Optional[float], optional): Momentum for running_mean or running_var
507                computation (used in batchnorm). If None, a cumulative moving average is used.
508                Defaults to 0.1.
509            affine (bool, optional): If True, the module has learnable affine parameters
510                (used in batchnorm). Defaults to True.
511            ratio (float, optional): Multiplier for the geometric progression of increasing channels
512                (feature maps). Used for 'channel_growth_rate' as 'exponential' or 'power'.
513                Defaults to 2 (powers of two).
514            channel_growth_rate (str, optional): Method of calculating the number of feature maps.
515                Options include 'exponential', 'proportion', 'linear', 'power', and 'constant'.
516                Defaults to 'exponential'.
517            conv_dim (int, optional): The dimension of the convolutional operation.
518                Defaults to 2.
519            adaptive_pool (str, optional): Type of adaptive pooling layer to apply last,
520                can be 'avgpool' or 'maxpool'. Defaults to 'avgpool'.
521
522        Returns:
523            A sequential block containing the specified transposed convolutional
524            layers, possibly including normalization and adaptive pooling.
525        """
526        _validate_input_size_is_not_none(self.input_size)
527        params = _set_conv_params(default_params=self._default_transpose_params, params=params)
528        conv_out = _select_conv_calc(conv_dim=conv_dim, transpose=True)
529
530        modules = []
531
532        if in_channels is None and self._conv_channels:
533            in_channels = self._conv_channels[-1]
534
535        _validate_build_transpose_convolve_init(in_channels, self._conv_channels)
536
537        if self._conv_layers:
538            input_layer_size_list = [self._conv_layers[-1]]
539
540        input_channels_count_list = self._calc_out_transpose_channels(
541            in_channels=in_channels,
542            out_channels=out_channels,
543            n_layers=n_layers,
544            ratio=ratio,
545            channel_growth_rate=channel_growth_rate,
546        )
547        for layer in range(n_layers):
548            _validate_max_channels_number(layer, input_channels_count_list, self.max_channels)
549            _validate_min_channels_number(layer, input_channels_count_list, min_channels=1)
550
551            in_channels = input_channels_count_list[layer]
552            out_channels = input_channels_count_list[layer + 1]
553
554            if self._conv_layers:
555                input_layer_size = input_layer_size_list[-1]
556                out_layer_size = conv_out(input_size=input_layer_size, **params)
557                input_layer_size_list.append(out_layer_size)
558
559            last_block_condition = layer == n_layers - 1
560            convolve_block = self.build_transpose_convolve_block(
561                in_channels=in_channels,
562                out_channels=out_channels,
563                normalization=normalization,
564                sub_blocks=sub_blocks,
565                p=p,
566                inplace=inplace,
567                eps=eps,
568                momentum=momentum,
569                affine=affine,
570                params=params,
571                last_block=last_block_condition,
572                conv_dim=conv_dim,
573            )
574
575            modules.append((f"deconv {layer + 1}", convolve_block))
576
577        self._transpose_conv_channels = input_channels_count_list
578
579        if self._conv_layers:
580            self._transpose_conv_layers = input_layer_size_list
581
582        if out_size is None:
583            out_size = self.input_size
584
585        adaptive_pooling = _select_adaptive_pooling_dimension(conv_dim=conv_dim, pooling=adaptive_pool)
586        resize_block = adaptive_pooling(output_size=tuple(out_size))
587        modules.append(("resize", resize_block))
588
589        return nn.Sequential(OrderedDict(modules))
590
591    def latent_block(
592        self,
593        input_shape: Sequence[int],
594        output_shape: Sequence[int],
595        n_layers: int = 1,
596        activation_function: Union[Optional[nn.Module], str] = None,
597    ):
598        """Creates a latent space transformation block.
599
600        This method constructs a latent space module that transforms an input tensor
601        of a specified shape into an output tensor of a desired shape using
602        linear layers and an optional activation function.
603
604        Args:
605            input_shape (Sequence[int]): The shape of the input tensor, typically
606                represented as a sequence of integers.
607            output_shape (Sequence[int]): The desired shape of the output tensor,
608                specified as a sequence of integers.
609            n_layers (int, optional): The number of linear layers to use in the
610                transformation. Defaults to 1.
611            activation_function (Union[Optional[nn.Module], str], optional):
612                Specifies the activation function to apply after the linear layers.
613                If set to 'same', it will use the instance's predefined activation
614                function. Defaults to None.
615
616        Returns:
617            An instance of the LatentSpaceModule class that
618            performs the specified transformation from input to output shape.
619        """
620        if activation_function == "same":
621            activation_function = self.activation_function
622
623        return LatentSpaceModule(input_shape, output_shape, n_layers, activation_function)
624
625    def _calc_out_channels(
626        self,
627        in_size: Sequence[int],
628        in_channels: int,
629        n_layers: int,
630        ratio: float = 2.0,
631        start: int = 32,
632        constant: int = 1,
633        channel_growth_rate: str = "exponential",
634    ) -> List[int]:
635        _validate_channel_growth_rate_param(channel_growth_rate)
636
637        if channel_growth_rate == "exponential":
638            self.max_channels = self._initial_max_channels
639            return [in_channels] + [int(start * ratio**i) for i in range(n_layers)]
640
641        if channel_growth_rate == "proportion":
642            range_start = in_channels
643            range_stop = int((sum(in_size) * 0.5) // len(in_size) + in_channels)
644            range_step = (range_stop - in_channels) // n_layers
645
646            _validate_range_step(range_step, n_layers)
647
648            channels = list(range(range_start, range_stop + 1, range_step))[: n_layers + 1]
649            self.max_channels = range_stop
650            return channels
651
652        if channel_growth_rate == "linear":
653            self.max_channels = self.min_channels + n_layers
654            return [in_channels] + [in_channels + i + 1 for i in range(n_layers)]
655
656        if channel_growth_rate == "constant":
657            self.max_channels = constant + 1
658            return [in_channels] + [constant for _ in range(n_layers)]
659
660        if channel_growth_rate == "power":
661            self.max_channels = self._initial_max_channels
662            return [in_channels] + [int((in_channels + i) ** ratio) for i in range(1, n_layers + 1)]
663
664    @staticmethod
665    def _calc_out_transpose_channels(
666        in_channels: int,
667        out_channels: int,
668        n_layers: int,
669        ratio: float = 2.0,
670        constant: int = 1,
671        channel_growth_rate: str = "exponential",
672    ) -> List[int]:
673        _validate_channel_growth_rate_param(channel_growth_rate)
674
675        if channel_growth_rate == "exponential":
676            return [int(in_channels / ratio**i) for i in range(n_layers)] + [out_channels]
677
678        if channel_growth_rate == "proportion":
679            channels = list(range(out_channels, in_channels, (in_channels - out_channels) // n_layers))[::-1]
680            channels = channels[:n_layers]
681            channels[-1] = out_channels
682            return [in_channels] + channels
683
684        if channel_growth_rate == "linear":
685            return [in_channels] + [in_channels - i for i in range(1, n_layers)] + [out_channels]
686
687        if channel_growth_rate == "constant":
688            return [in_channels] + [constant for _ in range(n_layers - 1)] + [out_channels]
689
690        if channel_growth_rate == "power":
691            return (
692                [in_channels] + [int(math.pow((n_layers - i + 1), ratio)) for i in range(1, n_layers)] + [out_channels]
693            )