Exploring Stronger Transformer Representation Learning for Occluded Person Re-Identification

The occluded person re-identification mainly confronts the challenges of incomplete body information and spatial misalignment. The existing occluded person ReID solutions can be classified into two categories, methods based on the feature alignment of visible region cues [21, 22, 23] and methods based on information completion [24, 25, 26].

The core idea of methods based on the feature alignment of visible region cues is to accurately locate the visible regions of person in the image, which only relies on the semantic information of visible region to align the feature representations of different body parts and match the pedestrians with the same identity in other camera view. Miao et al. [21] proposed a feature alignment method based on the human key-points information, taking advantage of human key-points information to guide the model to focus on the non-occluded regions and only these are utilized for the retrieval. He et al. [22] designed a deep spatial feature reconstruction method for partial person ReID, which realized the implicit feature alignment when calculating the reconstruction error of the spatial feature maps. However, solving the reconstruction error needs the time-consuming computation and is not suitable for the large scale re-identification task. Chen et al. [23] proposed an occlusion-aware mask network, which utilizes the attention-guided mask module to precisely capture body parts regardless of the occlusion. Wang et al. [27] proposed a novel feature erasing and diffusion network to simultaneously handle the interference from non-pedestrian occlusions and non-target pedestrians, improving the model’s perception ability towards target pedestrians and robustness towards non-target pedestrians.

The person ReID methods based on information completion usually adopt the spatiotemporal context to recover the pedestrian information of the occluded parts so as to improve the features’ discrimination [24, 26]. Xu et al. [25] proposed a feature recovery transformer to exploit the pedestrian information in its $k$-nearest neighbors features for occluded feature recovery. In [26], a spatiotemporal feature completion module was proposed to recover the occluded body parts via a region-encoder and a region-decoder for video person ReID.

Different from these methods above, we presented a joint-training model combining the supervised and self-supervised learning, which only introduces our specially designed data augmentation method in the training stage to learn more robust feature representation for the occlusion, not increasing the inference cost.

When capturing the long range dependencies of features, the CNN-based methods confront some challenges due to the limitation of receptive fields. In contrast, the ViT can overcome it by the self-attention mechanism, demonstrating excellent performance in some areas such as image understanding, image segmentation and visual question and answer. TransReID [4], as the first person ReID model to use the pure transformer architecture for extracting features, has further promoted the development of this field. The PFT [28] utilized the proposed three modules (patch full dimension enhancement module, fusion and reconstruction module and spatial slicing module) to enhance the efficiency of vision transformer. The PFD [29] utilized pose information to clearly disentangle semantic components (e.g. human body or joint parts) and selectively match non-occluded parts correspondingly, improving the performance of person ReID effectively. Li et al. [30] proposed a novel end-to-end part-aware transformer for occluded person ReID via a transformer encoder-decoder architecture, including a pixel context based transformer encoder and a part prototype based transformer decoder. In [6], the authors proposed the occlusion suppression and repairing transformer, including a self-supervised occlusion predictor, occlusion suppression encoder and feature repairing head, which can enhance the model’s ability of extracting discriminative local features for the occlusion scenes. There are other some methods to try to use the transformer to aggregate the features extracted from the CNN backbone, e.g., GiT [31], HaT [32] and so on.

Recently, self-supervised learning is widely applied to person re-identification task because it can mine the useful information from unlabelled data and has become an important direction of person ReID research, which can be classified into contrastive learning and masked image modeling.

Contrastive learning is aimed at learning similar/dissimilar representations from sample pairs with same semantic and ones with different semantic. The traditional contrastive learning usually relies on the annotated datasets where each image has a clear identity label. The core idea of this model is to attract the positive sample pairs and repulse the negative sample pairs by a reasonable designed loss function (e.g., Triplet loss [33], InfoNCE [34], contrastive loss [35] and so on) so as to learn the more discriminative features. Different from it, the self-supervised contrastive learning constructs positive and negative sample pairs from the strong and normal enhancing results of the same image, which can learn the meaningful features representation from lots of unlabelled data by the pre-designed tasks in the absence of labelled data. The existing self-supervised contrastive learning methods mainly include MoCo [36], SimCLR [19], BYOL [37], SimSiam [38] and so on.

Masked image modeling was first proposed in [39], which can enhance the ability of representation learning by performing pretext tasks such as image/feature reconstruction. Its basic idea is to randomly mask some image patches and reconstruct them from visible image patches by ViT in a self-supervised manner. This technology not only can be transferred to downstream computer vision tasks such as image classification and instance segmentation, but also provided a new way to solve the occluded person re-identification. In this paper, inspirited by the masked image modeling, we designed a novel data augmentation method for contrastive learning branch, that is random rectangle mask strategy.

Empirically, the occlusion often occurs for the person re-identification tasks. The normal data augmentation methods, e.g., random cropping, flipping and random color distortion, are not very effective in the face of occlusion. Thus, in order to improve the model’s performance under the situation of partial occlusion and incomplete information, inspirited by the paper [16], we designed a novel data augmentation method to imitate the occlusion in real scenes, that is random rectangle mask strategy, whose implementing method is given in detail in Algorithm 1 and the visualized results see the figure 1. This data augmentation method can help the model learn more discriminative features so as to facilitate effective recognition and classification in the case of incomplete information. In addition, in order to further enhance the model’s generalization, we also use other data augmentation methods, such as Gaussian blur, random color jittering and Solarization.

In [4], the TransReID has been proven to a well-performing person re-identification model based on the transformer. Based on it, we build our self-supervised and supervised combining transformer-based person re-identification framework, shown in Figure 2. Our model consists of the normal augmentation and strong augmentation branches, and they utilize the TransReID as the feature extractor and share the same weight of transformer layers. Specifically, each image is first processed by inputting the normal augmentation and strong augmentation branches respectively, where the normal augmentation includes some common methods, such as random horizontal flipping, random cropping, random erasing and so on, and the strong augmentation adopts the data augmentation methods mentioned in subsection III-A. Next, each augmented image is split into $N$ fixed-sized patches, and each patch is mapped to a $D$ dimensions vector by a linear projection. Subsequently, their patch embedding, position embedding and side information embedding are fed into $l-1$ layers transformer encoder together. Finally, the output features of the normal augmentation branch are fed into the global predicted head and jigsaw predicted head to further extract the features and calculate the ID loss and triplet loss, and the output features of the strong augmentation branch are input the projector and predictor in sequence to further process, obtaining features and the output ones of the normal branch calculate the similarity by the contrastive learning together. Our proposed model not only retains the Side Information Embedding (SIE) and jigsaw patch modules of the TransReID, but also assists features learning in the training phase by the contrastive learning branch, which greatly improves the model’s performance in complex scenes without affecting the efficiency of inference, especially for the occlusion.

Similar to the TransReID, we optimize the normal augmentation branch network by building ID loss and triplet loss for global features and $K$ fine-grained local features that obtained by the jigsaw patch module. The ID loss $\mathcal{L}_{ID}$ is the cross-entropy loss without label smoothing. For a triplet set $\{a,p,n\}$, the triplet loss $\mathcal{L}_{T}$ with soft-margin is defined as following

Thus, for the normal augmentation branch, the total loss of supervised learning can be represented as

where $f_{g}$ is served as the global feature obtained by the standard transformer encoder, and $f_{l}^{k}$ denotes the output token of the $k$-th group obtained by re-grouping them after shuffling the patches through the $l$-th transformer encoder.

In order to enhance the representation ability of the transformer encoder, we assist the training of ReID task by self-supervised contrastive learning. In the contrastive learning branch, we adopt the same loss function as the paper [38], which can be defined as following:

where $p_{1},p_{2}$ respectively denote the output vectors by the predictor of two stream branches, $z_{1},z_{2}$ denote the output features of the $l-1$-th transformer encoder of two stream architectures, $stopgrad(\cdot)$ is the stop-gradient operation which means that $z_{1},z_{2}$ are treated as the constant, that is, only receiving gradients from $p$ and no gradient from $z$ when encoding the corresponding input image, and $\mathcal{D}(\cdot,\cdot)$ is the negative cosine similarity.

Thus, the final total loss is defined as following:

where $\lambda$ is the hyper-parameter for balancing two losses. by adjusting it reasonably, we can let the contrastive loss assist the training of ReID task efficiently so as to utilize the advantage of self-supervised learning to enhance the performance and robustness of our proposed model against complex and variable scenes.

In this paper, all the experiments of our proposed method are run on a single NVIDIA A100 GPU with 40G RAM. All input person images are resize to $256\times 128$. For the normal augmentation branch, the training images are augmented with random horizontal flipping, random cropping and random erasing, and for another branch, they are augmented by the methods mentioned in subsection III-A. The maximum mask size $M$ is set to $(128,128)$ in Algorithm 1. The batch size is set to 100 with 4 images per ID. Stochastic Gradient Descent (SGD) optimizer is employed with a momentum of 0.9 and the weight decay of $1e-4$. The base learning rate is set to $0.0125$ with cosine learning rate decay. The projection head and prediction head are respectively a 3-layer MLP and 2-layer MLP, the hidden layers of which are 768-d and 4096-d and are with ReLU; the output layers are both 256-d, without ReLU. The hyper-parameter $\lambda$ is set to 0.95. The parameters of patch projecting layer in our model are randomly initialized and frozen rather than the learnable ones in the traditional methods. The initial weights of ViT adopt the officially released per-trained model, the model of which is first trained on Image-21K and then fine-tuned on ImageNet-1K.

In order to clearly clarify why our proposed method perform well on the occluded and general ReID datasets, we use Grad-CAM [53] to visualize the attention maps of the encoders from TransReID and our method by the heatmap, shown in Figure 3, which selects two occluded and one unoccluded images. Seeing from the Figure 3, comparing with the TransReID, our proposed method not only keeps paying attention to the overall features of target but also is not sensitive to the occlusion. For the occluded person, our proposed method makes the encoding features more focus on the unoccluded parts of person, which efficiently reduces the effect of the occlusion. This property of our proposed method ensures that it can extract more robust and discriminative features for person ReID task on the occluded and unoccluded images, which further verifies the superiority of our proposed method in person ReID task.

On Occluded-Duke dataset, we perform a number of ablation experiments to validate the effectiveness of our proposed method from different aspects.

For strong augmentation branch, we compare our random mask with four most similar data augmentation methods, the results of which are shown in Table V. Although these data augmentation methods are designed to simulate occlusion so as to improve the generalization ability of the model, we observe that our random mask and Hide-and-seek perform better than other three methods. The main reason is that our random mask and Hide-and-seek simulate more severe occlusion, compared with other three methods (the occlusion ratio of each sample is set to 0.5 for our random mask and Hide-and-seek, and the occlusion ratio of each sample is randomly selected from 0 to 0.5 for other three methods), which is beneficial for the contrastive learning. Seeing from the Table V, our random mask is superior to Hide-and-seek. Even though they have the same mAP score, the Rank-1 score of our random mask is 0.6% higher than the Rank-1 one of Hide-and-seek. The possible reason is that the occlusion simulated by our method is more diverse and complex, and the one generated by Hide-and-seek is uniformly and randomly distributed, because more complex and diverse samples are helpful for contrastive learning. In Table VI, we compare the results of our random mask in combination with other several data augmentation methods (Color Jitter, Gaussian Blur and Solarize), and we find that our Rank-1 score rises by 0.5% when combining simultaneously with those three augmentation methods.