Skip to content

Model

HerdNet

Bases: Module

HerdNet architecture

Source code in PytorchWildlife/models/detection/herdnet/model.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 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
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
class HerdNet(nn.Module):
    ''' HerdNet architecture '''

    def __init__(
        self,
        num_layers: int = 34,
        num_classes: int = 2,
        pretrained: bool = True, 
        down_ratio: Optional[int] = 2, 
        head_conv: int = 64
        ):
        '''
        Args:
            num_layers (int, optional): number of layers of DLA. Defaults to 34.
            num_classes (int, optional): number of output classes, background included. 
                Defaults to 2.
            pretrained (bool, optional): set False to disable pretrained DLA encoder parameters
                from ImageNet. Defaults to True.
            down_ratio (int, optional): downsample ratio. Possible values are 1, 2, 4, 8, or 16. 
                Set to 1 to get output of the same size as input (i.e. no downsample).
                Defaults to 2.
            head_conv (int, optional): number of supplementary convolutional layers at the end 
                of decoder. Defaults to 64.
        '''

        super(HerdNet, self).__init__()

        assert down_ratio in [1, 2, 4, 8, 16], \
            f'Downsample ratio possible values are 1, 2, 4, 8 or 16, got {down_ratio}'

        base_name = 'dla{}'.format(num_layers)

        self.down_ratio = down_ratio
        self.num_classes = num_classes
        self.head_conv = head_conv

        self.first_level = int(np.log2(down_ratio))

        # backbone
        base = dla_modules.__dict__[base_name](pretrained=pretrained, return_levels=True)
        setattr(self, 'base_0', base)
        setattr(self, 'channels_0', base.channels)

        channels = self.channels_0

        scales = [2 ** i for i in range(len(channels[self.first_level:]))]
        self.dla_up = dla_modules.DLAUp(channels[self.first_level:], scales=scales)
        # self.cls_dla_up = dla_modules.DLAUp(channels[-3:], scales=scales[:3])

        # bottleneck conv
        self.bottleneck_conv = nn.Conv2d(
            channels[-1], channels[-1], 
            kernel_size=1, stride=1, 
            padding=0, bias=True
        )

        # localization head
        self.loc_head = nn.Sequential(
            nn.Conv2d(channels[self.first_level], head_conv,
            kernel_size=3, padding=1, bias=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(
                head_conv, 1, 
                kernel_size=1, stride=1, 
                padding=0, bias=True
                ),
            nn.Sigmoid()
            )

        self.loc_head[-2].bias.data.fill_(0.00)

        # classification head
        self.cls_head = nn.Sequential(
            nn.Conv2d(channels[-1], head_conv,
            kernel_size=3, padding=1, bias=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(
                head_conv, self.num_classes, 
                kernel_size=1, stride=1, 
                padding=0, bias=True
                )
            )

        self.cls_head[-1].bias.data.fill_(0.00)

    def forward(self, input: torch.Tensor):

        encode = self.base_0(input)    
        bottleneck = self.bottleneck_conv(encode[-1])
        encode[-1] = bottleneck

        decode_hm = self.dla_up(encode[self.first_level:])
        # decode_cls = self.cls_dla_up(encode[-3:])

        heatmap = self.loc_head(decode_hm)
        clsmap = self.cls_head(bottleneck)
        # clsmap = self.cls_head(decode_cls)

        return heatmap, clsmap

    def freeze(self, layers: list) -> None:
        ''' Freeze all layers mentioned in the input list '''
        for layer in layers:
            self._freeze_layer(layer)

    def _freeze_layer(self, layer_name: str) -> None:
        for param in getattr(self, layer_name).parameters():
            param.requires_grad = False

    def reshape_classes(self, num_classes: int) -> None:
        ''' Reshape architecture according to a new number of classes.

        Arg:
            num_classes (int): new number of classes
        '''

        self.cls_head[-1] = nn.Conv2d(
                self.head_conv, num_classes, 
                kernel_size=1, stride=1, 
                padding=0, bias=True
                )

        self.cls_head[-1].bias.data.fill_(0.00)

        self.num_classes = num_classes

__init__(num_layers=34, num_classes=2, pretrained=True, down_ratio=2, head_conv=64)

Parameters:

Name Type Description Default
num_layers int

number of layers of DLA. Defaults to 34.

 34
num_classes int

number of output classes, background included. Defaults to 2.

 2
pretrained bool

set False to disable pretrained DLA encoder parameters from ImageNet. Defaults to True.

 True
down_ratio int

downsample ratio. Possible values are 1, 2, 4, 8, or 16. Set to 1 to get output of the same size as input (i.e. no downsample). Defaults to 2.

 2
head_conv int

number of supplementary convolutional layers at the end of decoder. Defaults to 64.

 64
Source code in PytorchWildlife/models/detection/herdnet/model.py 
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 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
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def __init__(
    self,
    num_layers: int = 34,
    num_classes: int = 2,
    pretrained: bool = True, 
    down_ratio: Optional[int] = 2, 
    head_conv: int = 64
    ):
    '''
    Args:
        num_layers (int, optional): number of layers of DLA. Defaults to 34.
        num_classes (int, optional): number of output classes, background included. 
            Defaults to 2.
        pretrained (bool, optional): set False to disable pretrained DLA encoder parameters
            from ImageNet. Defaults to True.
        down_ratio (int, optional): downsample ratio. Possible values are 1, 2, 4, 8, or 16. 
            Set to 1 to get output of the same size as input (i.e. no downsample).
            Defaults to 2.
        head_conv (int, optional): number of supplementary convolutional layers at the end 
            of decoder. Defaults to 64.
    '''

    super(HerdNet, self).__init__()

    assert down_ratio in [1, 2, 4, 8, 16], \
        f'Downsample ratio possible values are 1, 2, 4, 8 or 16, got {down_ratio}'

    base_name = 'dla{}'.format(num_layers)

    self.down_ratio = down_ratio
    self.num_classes = num_classes
    self.head_conv = head_conv

    self.first_level = int(np.log2(down_ratio))

    # backbone
    base = dla_modules.__dict__[base_name](pretrained=pretrained, return_levels=True)
    setattr(self, 'base_0', base)
    setattr(self, 'channels_0', base.channels)

    channels = self.channels_0

    scales = [2 ** i for i in range(len(channels[self.first_level:]))]
    self.dla_up = dla_modules.DLAUp(channels[self.first_level:], scales=scales)
    # self.cls_dla_up = dla_modules.DLAUp(channels[-3:], scales=scales[:3])

    # bottleneck conv
    self.bottleneck_conv = nn.Conv2d(
        channels[-1], channels[-1], 
        kernel_size=1, stride=1, 
        padding=0, bias=True
    )

    # localization head
    self.loc_head = nn.Sequential(
        nn.Conv2d(channels[self.first_level], head_conv,
        kernel_size=3, padding=1, bias=True),
        nn.ReLU(inplace=True),
        nn.Conv2d(
            head_conv, 1, 
            kernel_size=1, stride=1, 
            padding=0, bias=True
            ),
        nn.Sigmoid()
        )

    self.loc_head[-2].bias.data.fill_(0.00)

    # classification head
    self.cls_head = nn.Sequential(
        nn.Conv2d(channels[-1], head_conv,
        kernel_size=3, padding=1, bias=True),
        nn.ReLU(inplace=True),
        nn.Conv2d(
            head_conv, self.num_classes, 
            kernel_size=1, stride=1, 
            padding=0, bias=True
            )
        )

    self.cls_head[-1].bias.data.fill_(0.00)

freeze(layers)

Freeze all layers mentioned in the input list

Source code in PytorchWildlife/models/detection/herdnet/model.py 
127
128
129
130
def freeze(self, layers: list) -> None:
    ''' Freeze all layers mentioned in the input list '''
    for layer in layers:
        self._freeze_layer(layer)

reshape_classes(num_classes)

Reshape architecture according to a new number of classes.

Arg

num_classes (int): new number of classes

Source code in PytorchWildlife/models/detection/herdnet/model.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def reshape_classes(self, num_classes: int) -> None:
    ''' Reshape architecture according to a new number of classes.

    Arg:
        num_classes (int): new number of classes
    '''

    self.cls_head[-1] = nn.Conv2d(
            self.head_conv, num_classes, 
            kernel_size=1, stride=1, 
            padding=0, bias=True
            )

    self.cls_head[-1].bias.data.fill_(0.00)

    self.num_classes = num_classes