Skip to content

deeplabcut.pose_estimation_pytorch.models.weight_init

Ways to initialize weights for PyTorch modules.

Classes:

Name Description
BaseWeightInitializer

Class to used to initialize model weights.
Dekr

Class to used to initialize model weights in the same way as DEKR.
Normal

Class to used to initialize model weights using a normal distribution.
Rtmpose

Class to used to initialize head weights in the same way as RTMPose.

BaseWeightInitializer

Bases: ABC

Class to used to initialize model weights.

Methods:

Name Description
init_weights

Initializes weights for a model.
Source code in deeplabcut/pose_estimation_pytorch/models/weight_init.py 
30
31
32
33
34
35
36
37
38
39
class BaseWeightInitializer(ABC):
    """Class to used to initialize model weights."""

    @abstractmethod
    def init_weights(self, model: nn.Module) -> None:
        """Initializes weights for a model.

        Args:
            model: The model for which to initialize weights
        """

init_weights abstractmethod 

init_weights(model: Module) -> None

Initializes weights for a model.

Parameters:

Name Type Description Default

model

Module

The model for which to initialize weights

 required
Source code in deeplabcut/pose_estimation_pytorch/models/weight_init.py 
33
34
35
36
37
38
39
@abstractmethod
def init_weights(self, model: nn.Module) -> None:
    """Initializes weights for a model.

    Args:
        model: The model for which to initialize weights
    """

Dekr

Bases: BaseWeightInitializer

Class to used to initialize model weights in the same way as DEKR.

Attributes:

Name Type Description
std

the standard deviation to use to initialize weights
Source code in deeplabcut/pose_estimation_pytorch/models/weight_init.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
@WEIGHT_INIT.register_module
class Dekr(BaseWeightInitializer):
    """Class to used to initialize model weights in the same way as DEKR.

    Attributes:
        std: the standard deviation to use to initialize weights
    """

    def __init__(self, std: float = 0.001):
        self.std = std

    def init_weights(self, model: nn.Module) -> None:
        for name, module in model.named_parameters():
            if "bias" in name:
                nn.init.constant_(module, 0)
            else:
                nn.init.normal_(module, std=self.std)

            if hasattr(module, "transform_matrix_conv"):
                nn.init.constant_(module.transform_matrix_conv.weight, 0)
                if hasattr(module, "bias"):
                    nn.init.constant_(module.transform_matrix_conv.bias, 0)
            if hasattr(module, "translation_conv"):
                nn.init.constant_(module.translation_conv.weight, 0)
                if hasattr(module, "bias"):
                    nn.init.constant_(module.translation_conv.bias, 0)

Normal

Bases: BaseWeightInitializer

Class to used to initialize model weights using a normal distribution.

Weights are initialized with a normal distribution, and biases are initialized to 0.

Attributes:

Name Type Description
std

the standard deviation to use to initialize weights
Source code in deeplabcut/pose_estimation_pytorch/models/weight_init.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
@WEIGHT_INIT.register_module
class Normal(BaseWeightInitializer):
    """Class to used to initialize model weights using a normal distribution.

    Weights are initialized with a normal distribution, and biases are initialized to 0.

    Attributes:
        std: the standard deviation to use to initialize weights
    """

    def __init__(self, std: float = 0.001):
        self.std = std

    def init_weights(self, model: nn.Module) -> None:
        for name, module in model.named_parameters():
            if "bias" in name:
                nn.init.constant_(module, 0)
            else:
                nn.init.normal_(module, std=self.std)

Rtmpose

Bases: BaseWeightInitializer

Class to used to initialize head weights in the same way as RTMPose.

Source code in deeplabcut/pose_estimation_pytorch/models/weight_init.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@WEIGHT_INIT.register_module
class Rtmpose(BaseWeightInitializer):
    """Class to used to initialize head weights in the same way as RTMPose."""

    def init_weights(self, model: nn.Module) -> None:
        for module in model.modules():
            if isinstance(module, nn.Conv2d):
                nn.init.normal_(module.weight, std=0.001)
                nn.init.constant_(module.bias, 0)
            elif isinstance(module, nn.BatchNorm2d):
                nn.init.constant_(module.weight, 1)
                nn.init.constant_(module.bias, 1)
            elif isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, std=0.01)
                nn.init.constant_(module.bias, 0)