Support SUSTech1K

opengait/modeling/models/baseline.py

@@ -3,6 +3,7 @@ import torch
 from ..base_model import BaseModel
 from ..modules import SetBlockWrapper, HorizontalPoolingPyramid, PackSequenceWrapper, SeparateFCs, SeparateBNNecks
 
+from einops import rearrange
 
 class Baseline(BaseModel):
 
@@ -20,6 +21,8 @@ class Baseline(BaseModel):
         sils = ipts[0]
         if len(sils.size()) == 4:
             sils = sils.unsqueeze(1)
+        else:
+            sils = rearrange(sils, 'n s c h w -> n c s h w')
 
         del ipts
         outs = self.Backbone(sils)  # [n, c, s, h, w]
@@ -33,17 +36,16 @@ class Baseline(BaseModel):
         embed_2, logits = self.BNNecks(embed_1)  # [n, c, p]
         embed = embed_1
 
-        n, _, s, h, w = sils.size()
         retval = {
             'training_feat': {
                 'triplet': {'embeddings': embed_1, 'labels': labs},
                 'softmax': {'logits': logits, 'labels': labs}
             },
             'visual_summary': {
-                'image/sils': sils.view(n*s, 1, h, w)
+                'image/sils': rearrange(sils,'n c s h w -> (n s) c h w')
             },
             'inference_feat': {
                 'embeddings': embed
             }
         }
-        return retval
+        return retval