def create_resnet_basic_3d_module(
    *,
    in_channels,
    out_channels,
    conv_kernel_size=(3, 7, 7),
    conv_stride=(1, 2, 2),
    conv_padding=(1, 3, 3),
    conv_bias=False,
    conv=nn.Conv3d,
    pool=nn.MaxPool3d,
    pool_kernel_size=(1, 3, 3),
    pool_stride=(1, 2, 2),
    pool_padding=(0, 1, 1),
    norm=nn.BatchNorm3d,
    norm_eps=1e-5,
    norm_momentum=0.1,
    activation=nn.ReLU,
) -> nn.Module:
    """
    ResNet basic 3D stem module. It performs spatiotemporal Convolution, BN, and activation
    following by a spatiotemporal pooling.
                                        Conv3d
                                           ↓
                                     Normalization
                                           ↓
                                       Activation
                                           ↓
                                        Pool3d
    Args:
        in_channels (`int`):
            Input channel size of the convolution.
        out_channels (`int`):
            Output channel size of the convolution.
        conv_kernel_size (`Tuple`):
            Convolutional kernel size(s).
        conv_stride (`Tuple`):
            Convolutional stride size(s).
        conv_padding (`Tuple`):
            Convolutional padding size(s).
        conv_bias (`bool`):
            Convolutional bias. If true, adds a learnable bias to the output.
        conv (`callable`):
            Callable used to build the convolution layer.
        pool (`Callable`):
            A callable that constructs pooling layer, options include: nn.AvgPool3d,
            nn.MaxPool3d, and None (not performing pooling).
        pool_kernel_size (`Tuple`):
            Pooling kernel size(s).
        pool_stride (`Tuple`):
            Pooling stride size(s).
        pool_padding (`Tuple`):
            Pooling padding size(s).
        norm (`Callable`):
            A callable that constructs normalization layer, options
            include nn.BatchNorm3d, None (not performing normalization).
        norm_eps (`float`):
            Normalization epsilon.
        norm_momentum (`float`):
            Normalization momentum.
        activation (`Callable`):
            A callable that constructs activation layer, options
            include: nn.ReLU, nn.Softmax, nn.Sigmoid, and None (not performing activation).
    """

    assert conv is not None

    conv_module = conv(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=conv_kernel_size,
        stride=conv_stride,
        padding=conv_padding,
        bias=conv_bias,
    )
    norm_module = (
        None
        if norm is None
        else norm(num_features=out_channels, eps=norm_eps, momentum=norm_momentum)
    )
    activation_module = (
        None
        if activation is None
        else activation()
    )
    pool_module = (
        None
        if pool is None
        else pool(
            kernel_size=pool_kernel_size, stride=pool_stride, padding=pool_padding
        )
    )
    return ResNetBasic3DModule(
        conv=conv_module,
        norm=norm_module,
        activation=activation_module,
        pool=pool_module,
        )
ResNet basic 3D stem module. It performs spatiotemporal Convolution, BN, and activation
following by a spatiotemporal pooling.
Conv3d
↓
Normalization
↓
Activation
↓
Pool3d
Args:
in_channels (`int`):
Input channel size of the convolution.
out_channels (`int`):
Output channel size of the convolution.
conv_kernel_size (`Tuple`):
Convolutional kernel size(s).
conv_stride (`Tuple`):
Convolutional stride size(s).
conv_padding (`Tuple`):
Convolutional padding size(s).
conv_bias (`bool`):
Convolutional bias. If true, adds a learnable bias to the output.
conv (`callable`):
Callable used to build the convolution layer.
pool (`Callable`):
A callable that constructs pooling layer, options include: nn.AvgPool3d,
nn.MaxPool3d, and None (not performing pooling).
pool_kernel_size (`Tuple`):
Pooling kernel size(s).
pool_stride (`Tuple`):
Pooling stride size(s).
pool_padding (`Tuple`):
Pooling padding size(s).
norm (`Callable`):
A callable that constructs normalization layer, options
include nn.BatchNorm3d, None (not performing normalization).
norm_eps (`float`):
Normalization epsilon.
norm_momentum (`float`):
Normalization momentum.
activation (`Callable`):
A callable that constructs activation layer, options
include: nn.ReLU, nn.Softmax, nn.Sigmoid, and None (not performing activation).
towhee/models/utils/create_resnet_basic_3d_module.py