Add a flag to fix when dataset size is not divisible by batch size. (pytorch#1133)

Author: zuoxingdong

test/test_utils.py

@@ -8,6 +8,7 @@
 import unittest
 import traceback
 import torch
+import torch.utils.data
 import torch.cuda
 import warnings
 from torch.autograd import Variable
@@ -107,6 +108,44 @@ def __len__(self):
         return 10
 
 
+class TestDataLoader(TestCase):
+    def setUp(self):
+        self.dataset = torch.randn(5, 3, 3, 2)
+        self.batch_size = 3
+
+    def test_single_keep(self):
+        dataloader = torch.utils.data.DataLoader(self.dataset,
+                                                 batch_size=self.batch_size,
+                                                 num_workers=0,
+                                                 drop_last=False)
+        dataiter = iter(dataloader)
+        self.assertEqual(len(list(dataiter)), 2)
+
+    def test_single_drop(self):
+        dataloader = torch.utils.data.DataLoader(self.dataset,
+                                                 batch_size=self.batch_size,
+                                                 num_workers=0,
+                                                 drop_last=True)
+        dataiter = iter(dataloader)
+        self.assertEqual(len(list(dataiter)), 1)
+
+    def test_multi_keep(self):
+        dataloader = torch.utils.data.DataLoader(self.dataset,
+                                                 batch_size=self.batch_size,
+                                                 num_workers=2,
+                                                 drop_last=False)
+        dataiter = iter(dataloader)
+        self.assertEqual(len(list(dataiter)), 2)
+
+    def test_multi_drop(self):
+        dataloader = torch.utils.data.DataLoader(self.dataset,
+                                                 batch_size=self.batch_size,
+                                                 num_workers=2,
+                                                 drop_last=True)
+        dataiter = iter(dataloader)
+        self.assertEqual(len(list(dataiter)), 1)
+
+
 class TestTrainer(TestCase):
 
     intervals = [

torch/utils/data/dataloader.py

@@ -103,6 +103,7 @@ def __init__(self, loader):
         self.sampler = loader.sampler
         self.num_workers = loader.num_workers
         self.pin_memory = loader.pin_memory
+        self.drop_last = loader.drop_last
         self.done_event = threading.Event()
 
         self.samples_remaining = len(self.sampler)
@@ -141,11 +142,15 @@ def __init__(self, loader):
                 self._put_indices()
 
     def __len__(self):
-        return int(math.ceil(len(self.sampler) / float(self.batch_size)))
+        if self.drop_last:
+            return len(self.sampler) // self.batch_size
+        else:
+            return (len(self.sampler) + self.batch_size - 1) // self.batch_size
 
     def __next__(self):
-        if self.num_workers == 0:
-            # same-process loading
+        if self.num_workers == 0:  # same-process loading
+            if self.drop_last and self.samples_remaining < self.batch_size:
+                raise StopIteration
             if self.samples_remaining == 0:
                 raise StopIteration
             indices = self._next_indices()
@@ -187,9 +192,12 @@ def _next_indices(self):
     def _put_indices(self):
         assert self.batches_outstanding < 2 * self.num_workers
         if self.samples_remaining > 0:
-            self.index_queue.put((self.send_idx, self._next_indices()))
-            self.batches_outstanding += 1
-            self.send_idx += 1
+            if self.samples_remaining < self.batch_size and self.drop_last:
+                self._next_indices()
+            else:
+                self.index_queue.put((self.send_idx, self._next_indices()))
+                self.batches_outstanding += 1
+                self.send_idx += 1
 
     def _process_next_batch(self, batch):
         self.rcvd_idx += 1
@@ -236,15 +244,20 @@ class DataLoader(object):
             (default: 0)
         collate_fn (callable, optional)
         pin_memory (bool, optional)
+        drop_last (bool, optional): set to ``True`` to drop the last incomplete batch,
+            if the dataset size is not divisible by the batch size. If False and
+            the size of dataset is not divisible by the batch size, then the last batch
+            will be smaller. (default: False)
     """
 
-    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None,
-                 num_workers=0, collate_fn=default_collate, pin_memory=False):
+    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0,
+                 collate_fn=default_collate, pin_memory=False, drop_last=False):
         self.dataset = dataset
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.collate_fn = collate_fn
         self.pin_memory = pin_memory
+        self.drop_last = drop_last
 
         if sampler is not None:
             self.sampler = sampler
@@ -257,4 +270,7 @@ def __iter__(self):
         return DataLoaderIter(self)
 
     def __len__(self):
-        return int(math.ceil(len(self.sampler) / float(self.batch_size)))
+        if self.drop_last:
+            return len(self.sampler) // self.batch_size
+        else:
+            return (len(self.sampler) + self.batch_size - 1) // self.batch_size