File size: 5,813 Bytes
7b7527a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. 
#   
# Licensed under the Apache License, Version 2.0 (the "License");   
# you may not use this file except in compliance with the License.  
# You may obtain a copy of the License at   
#   
#     http://www.apache.org/licenses/LICENSE-2.0    
#   
# Unless required by applicable law or agreed to in writing, software   
# distributed under the License is distributed on an "AS IS" BASIS, 
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  
# See the License for the specific language governing permissions and   
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from paddle.nn.initializer import Normal, Constant
from paddle import ParamAttr
from .resnet import ResNet50, ResNet101
from ppdet.core.workspace import register

__all__ = ['PCBPyramid']


@register
class PCBPyramid(nn.Layer):
    """
    PCB (Part-based Convolutional Baseline), see https://arxiv.org/abs/1711.09349,
    Pyramidal Person Re-IDentification, see https://arxiv.org/abs/1810.12193

    Args:
        input_ch (int): Number of channels of the input feature.
        num_stripes (int): Number of sub-parts.
        used_levels (tuple): Whether the level is used, 1 means used.
        num_classes (int): Number of classes for identities, default 751 in
            Market-1501 dataset.
        last_conv_stride (int): Stride of the last conv.
        last_conv_dilation (int): Dilation of the last conv.
        num_conv_out_channels (int): Number of channels of conv feature.
    """

    def __init__(self,
                 input_ch=2048,
                 model_name='ResNet101',
                 num_stripes=6,
                 used_levels=(1, 1, 1, 1, 1, 1),
                 num_classes=751,
                 last_conv_stride=1,
                 last_conv_dilation=1,
                 num_conv_out_channels=128):
        super(PCBPyramid, self).__init__()
        self.num_stripes = num_stripes
        self.used_levels = used_levels
        self.num_classes = num_classes

        self.num_in_each_level = [i for i in range(self.num_stripes, 0, -1)]
        self.num_branches = sum(self.num_in_each_level)

        assert model_name in ['ResNet50', 'ResNet101'], "Unsupported ReID arch: {}".format(model_name)
        self.base = eval(model_name)(
            lr_mult=0.1,
            last_conv_stride=last_conv_stride,
            last_conv_dilation=last_conv_dilation)
        self.dropout_layer = nn.Dropout(p=0.2)
        self.pyramid_conv_list0, self.pyramid_fc_list0 = self.basic_branch(
            num_conv_out_channels, input_ch)

    def basic_branch(self, num_conv_out_channels, input_ch):
        # the level indexes are defined from fine to coarse,
        # the branch will contain one more part than that of its previous level
        # the sliding step is set to 1
        pyramid_conv_list = nn.LayerList()
        pyramid_fc_list = nn.LayerList()

        idx_levels = 0
        for idx_branches in range(self.num_branches):
            if idx_branches >= sum(self.num_in_each_level[0:idx_levels + 1]):
                idx_levels += 1

            pyramid_conv_list.append(
                nn.Sequential(
                    nn.Conv2D(input_ch, num_conv_out_channels, 1),
                    nn.BatchNorm2D(num_conv_out_channels), nn.ReLU()))

        idx_levels = 0
        for idx_branches in range(self.num_branches):
            if idx_branches >= sum(self.num_in_each_level[0:idx_levels + 1]):
                idx_levels += 1

            fc = nn.Linear(
                in_features=num_conv_out_channels,
                out_features=self.num_classes,
                weight_attr=ParamAttr(initializer=Normal(
                    mean=0., std=0.001)),
                bias_attr=ParamAttr(initializer=Constant(value=0.)))
            pyramid_fc_list.append(fc)
        return pyramid_conv_list, pyramid_fc_list

    def pyramid_forward(self, feat):
        each_stripe_size = int(feat.shape[2] / self.num_stripes)

        feat_list, logits_list = [], []
        idx_levels = 0
        used_branches = 0
        for idx_branches in range(self.num_branches):
            if idx_branches >= sum(self.num_in_each_level[0:idx_levels + 1]):
                idx_levels += 1
            idx_in_each_level = idx_branches - sum(self.num_in_each_level[
                0:idx_levels])
            stripe_size_in_each_level = each_stripe_size * (idx_levels + 1)
            start = idx_in_each_level * each_stripe_size
            end = start + stripe_size_in_each_level

            k = feat.shape[-1]
            local_feat_avgpool = F.avg_pool2d(
                feat[:, :, start:end, :],
                kernel_size=(stripe_size_in_each_level, k))
            local_feat_maxpool = F.max_pool2d(
                feat[:, :, start:end, :],
                kernel_size=(stripe_size_in_each_level, k))
            local_feat = local_feat_avgpool + local_feat_maxpool

            local_feat = self.pyramid_conv_list0[used_branches](local_feat)
            local_feat = paddle.reshape(
                local_feat, shape=[local_feat.shape[0], -1])
            feat_list.append(local_feat)

            local_logits = self.pyramid_fc_list0[used_branches](
                self.dropout_layer(local_feat))
            logits_list.append(local_logits)

            used_branches += 1

        return feat_list, logits_list

    def forward(self, x):
        feat = self.base(x)
        assert feat.shape[2] % self.num_stripes == 0
        feat_list, logits_list = self.pyramid_forward(feat)
        feat_out = paddle.concat(feat_list, axis=-1)
        return feat_out