From fa0728a0c7f7f7a47a93934d8f64317f1517fc15 Mon Sep 17 00:00:00 2001 From: Lingxiao He Date: Wed, 10 Jun 2020 17:43:56 +0800 Subject: [PATCH] dsr_head --- projects/PartialReID/partialreid/dsr_head.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/projects/PartialReID/partialreid/dsr_head.py b/projects/PartialReID/partialreid/dsr_head.py index 460f8fa..b90d4e6 100644 --- a/projects/PartialReID/partialreid/dsr_head.py +++ b/projects/PartialReID/partialreid/dsr_head.py @@ -35,8 +35,9 @@ class OcclusionUnit(nn.Module): SpatialFeatAll = SpatialFeatAll.transpose(1, 2) # shape: [n, c, m] y = self.mask_layer(SpatialFeatAll) mask_weight = torch.sigmoid(y[:, :, 0]) - - mask_score = F.normalize(mask_weight[:, :48], p=1, dim=1) + + feat_dim = SpaFeat1.size(2) * SpaFeat1.size(3) + mask_score = F.normalize(mask_weight[:, :feat_dim], p=1, dim=1) mask_weight_norm = F.normalize(mask_weight, p=1, dim=1) mask_score = mask_score.unsqueeze(1)