From 28c1afaa9d1fca38c296d86cecb8aef3ccd8449b Mon Sep 17 00:00:00 2001 From: Jeremy Reizenstein Date: Mon, 20 Jun 2022 13:47:47 -0700 Subject: [PATCH] nesting n_known_frames_for_test Summary: Use generator.permutation instead of choice so that different options for n_known_frames_for_test are nested. Reviewed By: davnov134 Differential Revision: D37210906 fbshipit-source-id: fd0d34ce62260417c3f63354a3f750aae9998b0d --- pytorch3d/implicitron/dataset/single_sequence_dataset.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pytorch3d/implicitron/dataset/single_sequence_dataset.py b/pytorch3d/implicitron/dataset/single_sequence_dataset.py index b52118ce..6f13f307 100644 --- a/pytorch3d/implicitron/dataset/single_sequence_dataset.py +++ b/pytorch3d/implicitron/dataset/single_sequence_dataset.py @@ -116,9 +116,11 @@ class SingleSceneDatasetMapProviderBase(DatasetMapProviderBase): if set_eval_batches: generator = np.random.default_rng(seed=0) for batch in eval_batches: - to_add = generator.choice( - len(train_split), self.n_known_frames_for_test - ) + # using permutation so that changes to n_known_frames_for_test + # result in consistent batches. + to_add = generator.permutation(len(train_split))[ + : self.n_known_frames_for_test + ] batch.extend((to_add + len(split)).tolist()) split = np.concatenate([split, train_split]) frame_types.extend([DATASET_TYPE_KNOWN] * len(train_split))