deeplabcut.pose_estimation_pytorch.models.modules.csp

Implementation of modules needed for the CSPNeXt Backbone. Used in CSP-style models.

Based on the building blocks used for the mmdetection CSPNeXt implementation. For more information, see https://github.com/open-mmlab/mmdetection.

Classes:

Name	Description
CSPConvModule	Configurable convolution module used for CSPNeXT.
CSPLayer	Cross Stage Partial Layer.
CSPNeXtBlock	Basic bottleneck block used in CSPNeXt.
ChannelAttention	Channel attention Module.
DepthwiseSeparableConv	Depth-wise separable convolution module used for CSPNeXT.
SPPBottleneck	Spatial pyramid pooling layer used in YOLOv3-SPP and (among others) CSPNeXt.

CSPConvModule

Bases: Module

Configurable convolution module used for CSPNeXT.

Applies sequentially - a convolution - (optional) a norm layer - (optional) an activation function

Parameters:

Name	Type	Description	Default
in_channels	int	Input channels of the convolution.	required
out_channels	int	Output channels of the convolution.	required
kernel_size	int | tuple[int, int]	Convolution kernel size.	required
stride	int | tuple[int, int]	Convolution stride.	1
padding	int | tuple[int, int]	Convolution padding.	0
dilation	int | tuple[int, int]	Convolution dilation.	1
groups	int	Number of blocked connections from input to output channels.	1
norm_layer	str | None	Norm layer to apply, if any.	None
activation_fn	str | None	Activation function to apply, if any.	'ReLU'

Source code in deeplabcut/pose_estimation_pytorch/models/modules/csp.py

class CSPConvModule(nn.Module):
    """Configurable convolution module used for CSPNeXT.

    Applies sequentially
      - a convolution
      - (optional) a norm layer
      - (optional) an activation function

    Args:
        in_channels: Input channels of the convolution.
        out_channels: Output channels of the convolution.
        kernel_size: Convolution kernel size.
        stride: Convolution stride.
        padding: Convolution padding.
        dilation: Convolution dilation.
        groups: Number of blocked connections from input to output channels.
        norm_layer: Norm layer to apply, if any.
        activation_fn: Activation function to apply, if any.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int | tuple[int, int],
        stride: int | tuple[int, int] = 1,
        padding: int | tuple[int, int] = 0,
        dilation: int | tuple[int, int] = 1,
        groups: int = 1,
        norm_layer: str | None = None,
        activation_fn: str | None = "ReLU",
    ):
        super().__init__()

        self.with_activation = activation_fn is not None
        self.with_bias = norm_layer is None
        self.with_norm = norm_layer is not None

        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=self.with_bias,
        )
        self.activate = None
        self.norm = None

        if self.with_norm:
            self.norm = build_norm(norm_layer, out_channels)

        if self.with_activation:
            # Careful when adding activation functions: some should not be in-place
            self.activate = build_activation(activation_fn, inplace=True)

        self._init_weights()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv(x)
        if self.with_norm:
            x = self.norm(x)
        if self.with_activation:
            x = self.activate(x)
        return x

    def _init_weights(self) -> None:
        """Same init as in <mmcv> convolutions."""
        nn.init.kaiming_normal_(self.conv.weight, a=0, nonlinearity="relu")
        if self.with_bias:
            nn.init.constant_(self.conv.bias, 0)

        if self.with_norm:
            nn.init.constant_(self.norm.weight, 1)
            nn.init.constant_(self.norm.bias, 0)

CSPLayer

Bases: Module

Cross Stage Partial Layer.

Parameters:

Name	Type	Description	Default
in_channels	int	input channels for the layer	required
out_channels	int	output channels for the block	required
expand_ratio	float	expansion factor for the mid-channels	0.5
num_blocks	int	the number of blocks to use	1
add_identity	bool	add a skip-connection to the blocks	True
channel_attention	bool	whether to apply channel attention	False
norm_layer	str | None	Norm layer to apply, if any.	None
activation_fn	str | None	Activation function to apply, if any.	'ReLU'

Methods:

Name	Description
forward	Forward function.

Source code in deeplabcut/pose_estimation_pytorch/models/modules/csp.py

class CSPLayer(nn.Module):
    """Cross Stage Partial Layer.

    Args:
        in_channels: input channels for the layer
        out_channels: output channels for the block
        expand_ratio: expansion factor for the mid-channels
        num_blocks: the number of blocks to use
        add_identity: add a skip-connection to the blocks
        channel_attention: whether to apply channel attention
        norm_layer: Norm layer to apply, if any.
        activation_fn: Activation function to apply, if any.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        expand_ratio: float = 0.5,
        num_blocks: int = 1,
        add_identity: bool = True,
        channel_attention: bool = False,
        norm_layer: str | None = None,
        activation_fn: str | None = "ReLU",
    ) -> None:
        super().__init__()
        mid_channels = int(out_channels * expand_ratio)
        self.channel_attention = channel_attention
        self.main_conv = CSPConvModule(
            in_channels,
            mid_channels,
            1,
            norm_layer=norm_layer,
            activation_fn=activation_fn,
        )
        self.short_conv = CSPConvModule(
            in_channels,
            mid_channels,
            1,
            norm_layer=norm_layer,
            activation_fn=activation_fn,
        )
        self.final_conv = CSPConvModule(
            2 * mid_channels,
            out_channels,
            1,
            norm_layer=norm_layer,
            activation_fn=activation_fn,
        )

        self.blocks = nn.Sequential(
            *[
                CSPNeXtBlock(
                    mid_channels,
                    mid_channels,
                    1.0,
                    add_identity,
                    norm_layer=norm_layer,
                    activation_fn=activation_fn,
                )
                for _ in range(num_blocks)
            ]
        )
        if channel_attention:
            self.attention = ChannelAttention(2 * mid_channels)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward function."""
        x_short = self.short_conv(x)

        x_main = self.main_conv(x)
        x_main = self.blocks(x_main)

        x_final = torch.cat((x_main, x_short), dim=1)

        if self.channel_attention:
            x_final = self.attention(x_final)
        return self.final_conv(x_final)

forward

forward(x: Tensor) -> torch.Tensor

Forward function.

Source code in deeplabcut/pose_estimation_pytorch/models/modules/csp.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward function."""
    x_short = self.short_conv(x)

    x_main = self.main_conv(x)
    x_main = self.blocks(x_main)

    x_final = torch.cat((x_main, x_short), dim=1)

    if self.channel_attention:
        x_final = self.attention(x_final)
    return self.final_conv(x_final)

CSPNeXtBlock

Bases: Module

Basic bottleneck block used in CSPNeXt.

Parameters:

Name	Type	Description	Default
in_channels	int	input channels for the block	required
out_channels	int	output channels for the block	required
expansion	float	expansion factor for the hidden channels	0.5
add_identity	bool	add a skip-connection to the block	True
kernel_size	int	kernel size for the DepthwiseSeparableConv	5
norm_layer	str | None	Norm layer to apply, if any.	None
activation_fn	str | None	Activation function to apply, if any.	'ReLU'

Methods:

Name	Description
forward	Forward function.

Source code in deeplabcut/pose_estimation_pytorch/models/modules/csp.py

class CSPNeXtBlock(nn.Module):
    """Basic bottleneck block used in CSPNeXt.

    Args:
        in_channels: input channels for the block
        out_channels: output channels for the block
        expansion: expansion factor for the hidden channels
        add_identity: add a skip-connection to the block
        kernel_size: kernel size for the DepthwiseSeparableConv
        norm_layer: Norm layer to apply, if any.
        activation_fn: Activation function to apply, if any.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        expansion: float = 0.5,
        add_identity: bool = True,
        kernel_size: int = 5,
        norm_layer: str | None = None,
        activation_fn: str | None = "ReLU",
    ) -> None:
        super().__init__()
        hidden_channels = int(out_channels * expansion)
        self.conv1 = CSPConvModule(
            in_channels,
            hidden_channels,
            3,
            stride=1,
            padding=1,
            norm_layer=norm_layer,
            activation_fn=activation_fn,
        )
        self.conv2 = DepthwiseSeparableConv(
            hidden_channels,
            out_channels,
            kernel_size,
            stride=1,
            padding=kernel_size // 2,
            norm_layer=norm_layer,
            activation_fn=activation_fn,
        )
        self.add_identity = add_identity and in_channels == out_channels

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward function."""
        identity = x
        out = self.conv1(x)
        out = self.conv2(out)

        if self.add_identity:
            return out + identity
        else:
            return out

forward

forward(x: Tensor) -> torch.Tensor

Forward function.

Source code in deeplabcut/pose_estimation_pytorch/models/modules/csp.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward function."""
    identity = x
    out = self.conv1(x)
    out = self.conv2(out)

    if self.add_identity:
        return out + identity
    else:
        return out

ChannelAttention

Bases: Module

Channel attention Module.

Parameters:

Name	Type	Description	Default
channels	int	Number of input/output channels of the layer.	required

Source code in deeplabcut/pose_estimation_pytorch/models/modules/csp.py

class ChannelAttention(nn.Module):
    """Channel attention Module.

    Args:
        channels: Number of input/output channels of the layer.
    """

    def __init__(self, channels: int) -> None:
        super().__init__()
        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
        self.act = nn.Hardsigmoid(inplace=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.amp.autocast("cuda", enabled=False):
            out = self.global_avgpool(x)
        out = self.fc(out)
        out = self.act(out)
        return x * out

DepthwiseSeparableConv

Bases: Module

Depth-wise separable convolution module used for CSPNeXT.

Applies sequentially - a depth-wise conv - a point-wise conv

Parameters:

Name	Type	Description	Default
in_channels	int	Input channels of the convolution.	required
out_channels	int	Output channels of the convolution.	required
kernel_size	int | tuple[int, int]	Convolution kernel size.	required
stride	int | tuple[int, int]	Convolution stride.	1
padding	int | tuple[int, int]	Convolution padding.	0
dilation	int | tuple[int, int]	Convolution dilation.	1
norm_layer	str | None	Norm layer to apply, if any.	None
activation_fn	str | None	Activation function to apply, if any.	'ReLU'

Source code in deeplabcut/pose_estimation_pytorch/models/modules/csp.py

class DepthwiseSeparableConv(nn.Module):
    """Depth-wise separable convolution module used for CSPNeXT.

    Applies sequentially
      - a depth-wise conv
      - a point-wise conv

    Args:
        in_channels: Input channels of the convolution.
        out_channels: Output channels of the convolution.
        kernel_size: Convolution kernel size.
        stride: Convolution stride.
        padding: Convolution padding.
        dilation: Convolution dilation.
        norm_layer: Norm layer to apply, if any.
        activation_fn: Activation function to apply, if any.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int | tuple[int, int],
        stride: int | tuple[int, int] = 1,
        padding: int | tuple[int, int] = 0,
        dilation: int | tuple[int, int] = 1,
        norm_layer: str | None = None,
        activation_fn: str | None = "ReLU",
    ):
        super().__init__()

        # depthwise convolution
        self.depthwise_conv = CSPConvModule(
            in_channels,
            in_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=in_channels,
            norm_layer=norm_layer,
            activation_fn=activation_fn,
        )

        self.pointwise_conv = CSPConvModule(
            in_channels,
            out_channels,
            kernel_size=1,
            norm_layer=norm_layer,
            activation_fn=activation_fn,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.depthwise_conv(x)
        x = self.pointwise_conv(x)
        return x

SPPBottleneck

Bases: Module

Spatial pyramid pooling layer used in YOLOv3-SPP and (among others) CSPNeXt.

Parameters:

Name	Type	Description	Default
in_channels	int	input channels to the bottleneck	required
out_channels	int	output channels of the bottleneck	required
kernel_sizes	tuple[int, ...]	kernel sizes for the pooling layers	(5, 9, 13)
norm_layer	str | None	norm layer for the bottleneck	'SyncBN'
activation_fn	str | None	activation function for the bottleneck	'SiLU'

Source code in deeplabcut/pose_estimation_pytorch/models/modules/csp.py

class SPPBottleneck(nn.Module):
    """Spatial pyramid pooling layer used in YOLOv3-SPP and (among others) CSPNeXt.

    Args:
        in_channels: input channels to the bottleneck
        out_channels: output channels of the bottleneck
        kernel_sizes: kernel sizes for the pooling layers
        norm_layer: norm layer for the bottleneck
        activation_fn: activation function for the bottleneck
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_sizes: tuple[int, ...] = (5, 9, 13),
        norm_layer: str | None = "SyncBN",
        activation_fn: str | None = "SiLU",
    ):
        super().__init__()
        mid_channels = in_channels // 2
        self.conv1 = CSPConvModule(
            in_channels,
            mid_channels,
            kernel_size=1,
            stride=1,
            norm_layer=norm_layer,
            activation_fn=activation_fn,
        )

        self.poolings = nn.ModuleList([nn.MaxPool2d(kernel_size=ks, stride=1, padding=ks // 2) for ks in kernel_sizes])
        conv2_channels = mid_channels * (len(kernel_sizes) + 1)
        self.conv2 = CSPConvModule(
            conv2_channels,
            out_channels,
            kernel_size=1,
            norm_layer=norm_layer,
            activation_fn=activation_fn,
        )

    def forward(self, x):
        x = self.conv1(x)
        with torch.amp.autocast("cuda", enabled=False):
            x = torch.cat([x] + [pooling(x) for pooling in self.poolings], dim=1)
        x = self.conv2(x)
        return x