fix(train): in case of last batch <=2, move to validation if possible

CHANGELOG.md

@@ -11,9 +11,8 @@ to [Semantic Versioning]. Full commit history is available in the
 #### Added
 
 - Added adaptive handling for last training minibatch of 1-2 cells in case of
-    `datasplitter_kwargs={"drop_last": False}` by moving them into validation set, if available.
+    `datasplitter_kwargs={"drop_last": False}` and `train_size = None` by moving them into validation set, if available.
     {pr}`3036`.
--
 
 #### Fixed
 

src/scvi/dataloaders/_data_splitting.py

@@ -25,8 +25,7 @@ def validate_data_split(
     n_samples: int,
     train_size: float,
     validation_size: float | None = None,
-    batch_size: int | None = None,
-    drop_last: bool | int = False,
+    batch_size_for_adaptive_last_batch: int | None = None,
 ):
     """Check data splitting parameters and return n_train and n_val.
 
@@ -38,10 +37,8 @@ def validate_data_split(
         Size of train set. Need to be: 0 < train_size <= 1.
     validation_size
         Size of validation set. Need to be 0 <= validation_size < 1
-    batch_size
-        batch size of each iteration. If `None`, do not minibatch
-    drop_last
-        drops last non-full batch
+    batch_size_for_adaptive_last_batch
+        batch size of each iteration. If `None`, do not do adaptive last batch sizing
     """
     if train_size > 1.0 or train_size <= 0.0:
         raise ValueError("Invalid train_size. Must be: 0 < train_size <= 1")
@@ -64,14 +61,15 @@ def validate_data_split(
             "any of the aforementioned parameters."
         )
 
-    if batch_size is not None and not drop_last:
-        if n_train % batch_size < 3 and n_train % batch_size > 0:
-            num_of_cells = n_train % batch_size
+    if batch_size_for_adaptive_last_batch is not None:
+        if (n_train % batch_size_for_adaptive_last_batch < 3 and
+            n_train % batch_size_for_adaptive_last_batch > 0):
+            num_of_cells = n_train % batch_size_for_adaptive_last_batch
             warnings.warn(
                 f"Last batch will have a small size of {num_of_cells} "
                 f"samples. Consider changing settings.batch_size or batch_size in model.train "
-                f"from currently {batch_size} to avoid errors during model training. Those cells "
-                f"will be removed from the training set automatically",
+                f"from currently {batch_size_for_adaptive_last_batch} to avoid errors during model"
+                f" training. Those cells will be removed from the training set automatically",
                 UserWarning,
                 stacklevel=settings.warnings_stacklevel,
             )
@@ -90,8 +88,6 @@ def validate_data_split(
 def validate_data_split_with_external_indexing(
     n_samples: int,
     external_indexing: list[np.array, np.array, np.array] | None = None,
-    batch_size: int | None = None,
-    drop_last: bool | int = False,
 ):
     """Check data splitting parameters and return n_train and n_val.
 
@@ -102,10 +98,6 @@ def validate_data_split_with_external_indexing(
     external_indexing
         A list of data split indices in the order of training, validation, and test sets.
         Validation and test set are not required and can be left empty.
-    batch_size
-        batch size of each iteration. If `None`, do not minibatch
-    drop_last
-        drops last non-full batch
     """
     if not isinstance(external_indexing, list):
         raise ValueError("External indexing is not of list type")
@@ -159,17 +151,6 @@ def validate_data_split_with_external_indexing(
     n_train = len(external_indexing[0])
     n_val = len(external_indexing[1])
 
-    if batch_size is not None and not drop_last:
-        if n_train % batch_size < 3 and n_train % batch_size > 0:
-            warnings.warn(
-                f"Last batch will have a small size of {n_train % batch_size} "
-                f"samples. Consider changing settings.batch_size or batch_size in model.train "
-                f"from currently {settings.batch_size} to avoid errors during model training "
-                f"or change the given external indices accordingly",
-                UserWarning,
-                stacklevel=settings.warnings_stacklevel,
-            )
-
     return n_train, n_val
 
 
@@ -183,7 +164,7 @@ class DataSplitter(pl.LightningDataModule):
     adata_manager
         :class:`~scvi.data.AnnDataManager` object that has been created via ``setup_anndata``.
     train_size
-        float, or None (default is 0.9)
+        float, or None (default is None, which is practicaly 0.9 + adaptive last batch)
     validation_size
         float, or None (default is None)
     shuffle_set_split
@@ -220,7 +201,7 @@ class is :class:`~scvi.dataloaders.SemiSupervisedDataLoader`,
     def __init__(
         self,
         adata_manager: AnnDataManager,
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         load_sparse_tensor: bool = False,
@@ -230,7 +211,8 @@ def __init__(
     ):
         super().__init__()
         self.adata_manager = adata_manager
-        self.train_size = float(train_size)
+        self.train_size_was_none = not bool(train_size)
+        self.train_size = 0.9 if self.train_size_was_none else float(train_size)
         self.validation_size = validation_size
         self.shuffle_set_split = shuffle_set_split
         self.load_sparse_tensor = load_sparse_tensor
@@ -243,16 +225,14 @@ def __init__(
             self.n_train, self.n_val = validate_data_split_with_external_indexing(
                 self.adata_manager.adata.n_obs,
                 self.external_indexing,
-                self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                self.drop_last,
             )
         else:
             self.n_train, self.n_val = validate_data_split(
                 self.adata_manager.adata.n_obs,
                 self.train_size,
                 self.validation_size,
-                self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                self.drop_last,
+                self.data_loader_kwargs.pop("batch_size", settings.batch_size) if
+                self.train_size_was_none and not self.drop_last else None,
             )
 
     def setup(self, stage: str | None = None):
@@ -342,7 +322,7 @@ class SemiSupervisedDataSplitter(pl.LightningDataModule):
     adata_manager
         :class:`~scvi.data.AnnDataManager` object that has been created via ``setup_anndata``.
     train_size
-        float, or None (default is 0.9)
+        float, or None (default is None, which is practicaly 0.9 + adaptive last batch)
     validation_size
         float, or None (default is None)
     shuffle_set_split
@@ -377,7 +357,7 @@ class is :class:`~scvi.dataloaders.SemiSupervisedDataLoader`,
     def __init__(
         self,
         adata_manager: AnnDataManager,
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         n_samples_per_label: int | None = None,
@@ -387,7 +367,8 @@ def __init__(
     ):
         super().__init__()
         self.adata_manager = adata_manager
-        self.train_size = float(train_size)
+        self.train_size_was_none = not bool(train_size)
+        self.train_size = 0.9 if train_size is None else float(train_size)
         self.validation_size = validation_size
         self.shuffle_set_split = shuffle_set_split
         self.drop_last = kwargs.pop("drop_last", False)
@@ -423,16 +404,14 @@ def setup(self, stage: str | None = None):
                 n_labeled_train, n_labeled_val = validate_data_split_with_external_indexing(
                     n_labeled_idx,
                     [labeled_idx_train, labeled_idx_val, labeled_idx_test],
-                    self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                    self.drop_last,
                 )
             else:
                 n_labeled_train, n_labeled_val = validate_data_split(
                     n_labeled_idx,
                     self.train_size,
                     self.validation_size,
-                    self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                    self.drop_last,
+                    self.data_loader_kwargs.pop("batch_size", settings.batch_size) if
+                    self.train_size_was_none and not self.drop_last else None,
                 )
 
                 labeled_permutation = self._labeled_indices
@@ -463,16 +442,14 @@ def setup(self, stage: str | None = None):
                 n_unlabeled_train, n_unlabeled_val = validate_data_split_with_external_indexing(
                     n_unlabeled_idx,
                     [unlabeled_idx_train, unlabeled_idx_val, unlabeled_idx_test],
-                    self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                    self.drop_last,
                 )
             else:
                 n_unlabeled_train, n_unlabeled_val = validate_data_split(
                     n_unlabeled_idx,
                     self.train_size,
                     self.validation_size,
-                    self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                    self.drop_last,
+                    self.data_loader_kwargs.pop("batch_size",settings.batch_size) if
+                    self.train_size_was_none and not self.drop_last else None,
                 )
 
                 unlabeled_permutation = self._unlabeled_indices
@@ -564,7 +541,7 @@ class DeviceBackedDataSplitter(DataSplitter):
     adata_manager
         :class:`~scvi.data.AnnDataManager` object that has been created via ``setup_anndata``.
     train_size
-        float, or None (default is 0.9)
+        float, or None (default is None, which is practicaly 0.9 + adaptive last batch)
     validation_size
         float, or None (default is None)
     %(param_accelerator)s
@@ -592,7 +569,7 @@ class DeviceBackedDataSplitter(DataSplitter):
     def __init__(
         self,
         adata_manager: AnnDataManager,
-        train_size: float = 1.0,
+        train_size: float | None = None,
         validation_size: float | None = None,
         accelerator: str = "auto",
         device: int | str = "auto",

src/scvi/external/cellassign/_model.py

@@ -143,7 +143,7 @@ def train(
         lr: float = 3e-3,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 1024,

src/scvi/external/contrastivevi/_contrastive_data_splitting.py

@@ -53,7 +53,7 @@ def __init__(
         adata_manager: AnnDataManager,
         background_indices: list[int],
         target_indices: list[int],
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         load_sparse_tensor: bool = False,
@@ -81,15 +81,15 @@ def __init__(
                 self.n_background,
                 self.train_size,
                 self.validation_size,
-                self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                self.drop_last,
+                self.data_loader_kwargs.pop("batch_size", settings.batch_size) if
+                self.train_size_was_none and not self.drop_last else None,
             )
             self.n_target_train, self.n_target_val = validate_data_split(
                 self.n_target,
                 self.train_size,
                 self.validation_size,
-                self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                self.drop_last,
+                self.data_loader_kwargs.pop("batch_size", settings.batch_size) if
+                self.train_size_was_none and not self.drop_last else None,
             )
         else:
             # we need to intersect the external indexing given with the bg/target indices
@@ -101,8 +101,6 @@ def __init__(
                 validate_data_split_with_external_indexing(
                     self.n_background,
                     [self.background_train_idx, self.background_val_idx, self.background_test_idx],
-                    self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                    self.drop_last,
                 )
             )
             self.background_train_idx, self.background_val_idx, self.background_test_idx = (
@@ -117,8 +115,6 @@ def __init__(
             self.n_target_train, self.n_target_val = validate_data_split_with_external_indexing(
                 self.n_target,
                 [self.target_train_idx, self.target_val_idx, self.target_test_idx],
-                self.data_loader_kwargs.pop("batch_size", settings.batch_size),
-                self.drop_last,
             )
             self.target_train_idx, self.target_val_idx, self.target_test_idx = (
                 self.target_train_idx.tolist(),

src/scvi/external/contrastivevi/_model.py

@@ -136,7 +136,7 @@ def train(
         max_epochs: int | None = None,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         load_sparse_tensor: bool = False,

src/scvi/external/gimvi/_model.py

@@ -172,7 +172,7 @@ def train(
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
         kappa: int = 5,
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 128,

src/scvi/external/mrvi/_model.py

@@ -201,7 +201,7 @@ def train(
         max_epochs: int | None = None,
         accelerator: str | None = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         batch_size: int = 128,
         early_stopping: bool = False,

src/scvi/external/scbasset/_model.py

@@ -106,7 +106,7 @@ def train(
         lr: float = 0.01,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 128,

src/scvi/external/solo/_model.py

@@ -293,7 +293,7 @@ def train(
         lr: float = 1e-3,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 128,

src/scvi/external/velovi/_model.py

@@ -129,7 +129,7 @@ def train(
         weight_decay: float = 1e-2,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         batch_size: int = 256,
         early_stopping: bool = True,

src/scvi/model/_multivi.py

@@ -231,7 +231,7 @@ def train(
         lr: float = 1e-4,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 128,

src/scvi/model/_peakvi.py

@@ -151,7 +151,7 @@ def train(
         lr: float = 1e-4,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 128,

src/scvi/model/_scanvi.py

@@ -356,7 +356,7 @@ def train(
         max_epochs: int | None = None,
         n_samples_per_label: float | None = None,
         check_val_every_n_epoch: int | None = None,
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 128,

src/scvi/model/_totalvi.py

@@ -197,7 +197,7 @@ def train(
         lr: float = 4e-3,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 256,

src/scvi/model/base/_jaxmixin.py

@@ -24,7 +24,7 @@ def train(
         max_epochs: int | None = None,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 128,

src/scvi/model/base/_pyromixin.py

@@ -92,7 +92,7 @@ def train(
         max_epochs: int | None = None,
         accelerator: str = "auto",
         device: int | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         batch_size: int = 128,

src/scvi/model/base/_training_mixin.py

@@ -24,7 +24,7 @@ def train(
         max_epochs: int | None = None,
         accelerator: str = "auto",
         devices: int | list[int] | str = "auto",
-        train_size: float = 0.9,
+        train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,
         load_sparse_tensor: bool = False,