ensure that SortishSampler always puts the largest sample first

stmsy · Apr 29, 2018 · 9f21874 · 9f21874
1 parent f29e102
commit 9f21874
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/fastai/text.py b/fastai/text.py
@@ -140,6 +140,8 @@ def __iter__(self):
         sort_idx = sum([sorted(s, key=self.key, reverse=True) for s in ck_idx], [])
         sz = self.bs
         ck_idx = [sort_idx[i:i+sz] for i in range(0, len(sort_idx), sz)]
+        max_ck = np.argmax([ck[0] for ck in ck_idx])
+        ck_idx[0],ck_idx[max_ck] = ck_idx[max_ck],ck_idx[0]
         sort_idx = np.concatenate(np.random.permutation(ck_idx[1:]))
         sort_idx = np.concatenate((ck_idx[0], sort_idx))
         return iter(sort_idx)

diff --git a/tests/test_samplers.py b/tests/test_samplers.py
@@ -28,5 +28,4 @@ def test_sortish_sampler_sorts_each_batch_descending():
         s1 > s2 or (i+1) % bs == 0  # don't check batch boundaries
         for i, (s1, s2) in enumerate(zip(samp, samp[1:]))
     )
-    # Not always true, though the class comment implies it should be.
-    # assert samp[0] == max(samp)
+    assert samp[0] == max(samp)