use dict as return of dataloader

Author: KaiyangZhou

torchreid/data/datasets/dataset.py

@@ -265,7 +265,13 @@ def __getitem__(self, index):
         img = read_image(img_path)
         if self.transform is not None:
             img = self.transform(img)
-        return img, pid, camid, img_path
+        item = {
+            'img': img,
+            'pid': pid,
+            'camid': camid,
+            'impath': img_path
+        }
+        return item
 
     def show_summary(self):
         num_train_pids, num_train_cams = self.parse_data(self.train)
@@ -373,7 +379,13 @@ def __getitem__(self, index):
             imgs.append(img)
         imgs = torch.cat(imgs, dim=0)
 
-        return imgs, pid, camid
+        item = {
+            'img': imgs,
+            'pid': pid,
+            'camid': camid
+        }
+
+        return item
 
     def show_summary(self):
         num_train_pids, num_train_cams = self.parse_data(self.train)

torchreid/engine/engine.py

@@ -449,14 +449,14 @@ def extract_features(self, input):
         return self.model(input)
 
     def parse_data_for_train(self, data):
-        imgs = data[0]
-        pids = data[1]
+        imgs = data['img']
+        pids = data['pid']
         return imgs, pids
 
     def parse_data_for_eval(self, data):
-        imgs = data[0]
-        pids = data[1]
-        camids = data[2]
+        imgs = data['img']
+        pids = data['pid']
+        camids = data['camid']
         return imgs, pids, camids
 
     def two_stepped_transfer_learning(

torchreid/engine/video/softmax.py

@@ -78,8 +78,8 @@ def __init__(
         self.pooling_method = pooling_method
 
     def parse_data_for_train(self, data):
-        imgs = data[0]
-        pids = data[1]
+        imgs = data['img']
+        pids = data['pid']
         if imgs.dim() == 5:
             # b: batch size
             # s: sqeuence length