fix bugs of dataloader

Author: xtt

dataset.py

@@ -36,7 +36,7 @@ def __init__(self, odgt, opt, max_sample=-1, batch_per_gpu=1):
 
         self.list_sample = [json.loads(x.rstrip()) for x in open(odgt, 'r')]
 
-        np.random.shuffle(self.list_sample)
+        self.if_shuffled = False
         if max_sample > 0:
             self.list_sample = self.list_sample[0:max_sample]
         self.num_sample = len(self.list_sample)
@@ -69,6 +69,11 @@ def _get_sub_batch(self):
         return batch_records
 
     def __getitem__(self, index):
+        # NOTE: random shuffle for the first time. shuffle in __init__ is useless
+        if not self.if_shuffled:
+            np.random.shuffle(self.list_sample)
+            self.if_shuffled = True
+
         # get sub-batch candidates
         batch_records = self._get_sub_batch()
 

train.py

@@ -52,8 +52,8 @@ def train(segmentation_module, iterator, optimizers, history, epoch, args):
         tic = time.time()
 
         # update average loss and acc
-        ave_total_loss.update(loss.item())
-        ave_acc.update(acc.item()*100)
+        ave_total_loss.update(loss.data[0])
+        ave_acc.update(acc.data[0]*100)
 
         # calculate accuracy, and display
         if i % args.disp_iter == 0:
@@ -67,8 +67,8 @@ def train(segmentation_module, iterator, optimizers, history, epoch, args):
 
             fractional_epoch = epoch - 1 + 1. * i / args.epoch_iters
             history['train']['epoch'].append(fractional_epoch)
-            history['train']['loss'].append(loss.item())
-            history['train']['acc'].append(acc.item())
+            history['train']['loss'].append(loss.data[0])
+            history['train']['acc'].append(acc.data[0])
 
         # adjust learning rate
         cur_iter = i + (epoch - 1) * args.epoch_iters
@@ -157,6 +157,8 @@ def main(args):
 
     # create loader iterator
     iterator_train = iter(loader_train)
+    from IPython import embed
+    embed()
 
     # load nets into gpu
     if args.num_gpus > 1: