deeplabcut.pose_estimation_pytorch.models.modules.conv_module

The code is based on DEKR: https://github.com/HRNet/DEKR/tree/main

Classes:

Name	Description
HighResolutionModule	High-Resolution Module.

HighResolutionModule

Bases: Module

High-Resolution Module.

This class implements the High-Resolution Module used in HigherHRNet for Human Pose Estimation.

Parameters:

Name	Type	Description	Default
num_branches	int	Number of branches in the module.	required
block	BasicBlock	The block type used in each branch of the module.	required
num_blocks	int	List containing the number of blocks in each branch.	required
num_inchannels	int	List containing the number of input channels for each branch.	required
num_channels	int	List containing the number of output channels for each branch.	required
fuse_method	str	The fusion method used in the module.	required
multi_scale_output	bool	Whether to output multi-scale features. Default is True.	True

Methods:

Name	Description
forward	Forward pass through the HighResolutionModule.

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

class HighResolutionModule(nn.Module):
    """High-Resolution Module.

    This class implements the High-Resolution Module used in HigherHRNet for Human Pose Estimation.

    Args:
        num_branches: Number of branches in the module.
        block: The block type used in each branch of the module.
        num_blocks: List containing the number of blocks in each branch.
        num_inchannels: List containing the number of input channels for each branch.
        num_channels: List containing the number of output channels for each branch.
        fuse_method: The fusion method used in the module.
        multi_scale_output: Whether to output multi-scale features. Default is True.
    """

    def __init__(
        self,
        num_branches: int,
        block: BasicBlock,
        num_blocks: int,
        num_inchannels: int,
        num_channels: int,
        fuse_method: str,
        multi_scale_output: bool = True,
    ):
        super().__init__()
        self._check_branches(num_branches, block, num_blocks, num_inchannels, num_channels)

        self.num_inchannels = num_inchannels
        self.fuse_method = fuse_method
        self.num_branches = num_branches

        self.multi_scale_output = multi_scale_output

        self.branches = self._make_branches(num_branches, block, num_blocks, num_channels)
        self.fuse_layers = self._make_fuse_layers()
        self.relu = nn.ReLU(True)

    def _check_branches(
        self,
        num_branches: int,
        block: BasicBlock,
        num_blocks: int,
        num_inchannels: int,
        num_channels: int,
    ):
        if num_branches != len(num_blocks):
            error_msg = f"NUM_BRANCHES({num_branches}) <> NUM_BLOCKS({len(num_blocks)})"
            logger.error(error_msg)
            raise ValueError(error_msg)

        if num_branches != len(num_channels):
            error_msg = f"NUM_BRANCHES({num_branches}) <> NUM_CHANNELS({len(num_channels)})"
            logger.error(error_msg)
            raise ValueError(error_msg)

        if num_branches != len(num_inchannels):
            error_msg = f"NUM_BRANCHES({num_branches}) <> NUM_INCHANNELS({len(num_inchannels)})"
            logger.error(error_msg)
            raise ValueError(error_msg)

    def _make_one_branch(
        self,
        branch_index: int,
        block: BasicBlock,
        num_blocks: int,
        num_channels: int,
        stride: int = 1,
    ) -> nn.Sequential:
        downsample = None
        if stride != 1 or self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(
                    self.num_inchannels[branch_index],
                    num_channels[branch_index] * block.expansion,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm2d(num_channels[branch_index] * block.expansion, momentum=BN_MOMENTUM),
            )

        layers = []
        layers.append(
            block(
                self.num_inchannels[branch_index],
                num_channels[branch_index],
                stride,
                downsample,
            )
        )
        self.num_inchannels[branch_index] = num_channels[branch_index] * block.expansion
        for _i in range(1, num_blocks[branch_index]):
            layers.append(block(self.num_inchannels[branch_index], num_channels[branch_index]))

        return nn.Sequential(*layers)

    def _make_branches(self, num_branches: int, block: BasicBlock, num_blocks: int, num_channels: int) -> nn.ModuleList:
        branches = []

        for i in range(num_branches):
            branches.append(self._make_one_branch(i, block, num_blocks, num_channels))

        return nn.ModuleList(branches)

    def _make_fuse_layers(self) -> nn.ModuleList:
        if self.num_branches == 1:
            return None

        num_branches = self.num_branches
        num_inchannels = self.num_inchannels
        fuse_layers = []
        for i in range(num_branches if self.multi_scale_output else 1):
            fuse_layer = []
            for j in range(num_branches):
                if j > i:
                    fuse_layer.append(
                        nn.Sequential(
                            nn.Conv2d(
                                num_inchannels[j],
                                num_inchannels[i],
                                1,
                                1,
                                0,
                                bias=False,
                            ),
                            nn.BatchNorm2d(num_inchannels[i]),
                            nn.Upsample(scale_factor=2 ** (j - i), mode="nearest"),
                        )
                    )
                elif j == i:
                    fuse_layer.append(None)
                else:
                    conv3x3s = []
                    for k in range(i - j):
                        if k == i - j - 1:
                            num_outchannels_conv3x3 = num_inchannels[i]
                            conv3x3s.append(
                                nn.Sequential(
                                    nn.Conv2d(
                                        num_inchannels[j],
                                        num_outchannels_conv3x3,
                                        3,
                                        2,
                                        1,
                                        bias=False,
                                    ),
                                    nn.BatchNorm2d(num_outchannels_conv3x3),
                                )
                            )
                        else:
                            num_outchannels_conv3x3 = num_inchannels[j]
                            conv3x3s.append(
                                nn.Sequential(
                                    nn.Conv2d(
                                        num_inchannels[j],
                                        num_outchannels_conv3x3,
                                        3,
                                        2,
                                        1,
                                        bias=False,
                                    ),
                                    nn.BatchNorm2d(num_outchannels_conv3x3),
                                    nn.ReLU(True),
                                )
                            )
                    fuse_layer.append(nn.Sequential(*conv3x3s))
            fuse_layers.append(nn.ModuleList(fuse_layer))

        return nn.ModuleList(fuse_layers)

    def get_num_inchannels(self) -> int:
        return self.num_inchannels

    def forward(self, x) -> list:
        """Forward pass through the HighResolutionModule.

        Args:
            x: List of input tensors for each branch.

        Returns:
            List of output tensors after processing through the module.
        """
        if self.num_branches == 1:
            return [self.branches[0](x[0])]

        for i in range(self.num_branches):
            x[i] = self.branches[i](x[i])

        x_fuse = []

        for i in range(len(self.fuse_layers)):
            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
            for j in range(1, self.num_branches):
                if i == j:
                    y = y + x[j]
                else:
                    y = y + self.fuse_layers[i][j](x[j])
            x_fuse.append(self.relu(y))

        return x_fuse

forward

forward(x) -> list

Forward pass through the HighResolutionModule.

Parameters:

Name	Type	Description	Default
x		List of input tensors for each branch.	required

Returns:

Type	Description
list	List of output tensors after processing through the module.

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

def forward(self, x) -> list:
    """Forward pass through the HighResolutionModule.

    Args:
        x: List of input tensors for each branch.

    Returns:
        List of output tensors after processing through the module.
    """
    if self.num_branches == 1:
        return [self.branches[0](x[0])]

    for i in range(self.num_branches):
        x[i] = self.branches[i](x[i])

    x_fuse = []

    for i in range(len(self.fuse_layers)):
        y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
        for j in range(1, self.num_branches):
            if i == j:
                y = y + x[j]
            else:
                y = y + self.fuse_layers[i][j](x[j])
        x_fuse.append(self.relu(y))

    return x_fuse