humancompatible-train 0.1.0__tar.gz → 0.1.2__tar.gz

{humancompatible_train-0.1.0/humancompatible_train.egg-info → humancompatible_train-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: humancompatible-train
-Version: 0.1.0
+Version: 0.1.2
 Summary: PyTorch-based package for constrained training of neural networks
 Author: Gilles Bareilles, Jana Lepsova, Jakub Marecek
 Author-email: Andrii Kliachkin <kliachkin.andrii@gmail.com>
@@ -10,30 +10,95 @@ Description-Content-Type: text/markdown
 License-File: LICENCE.txt
 Requires-Dist: torch
 Requires-Dist: numpy
+Provides-Extra: benchmark
+Requires-Dist: fairret; extra == "benchmark"
+Requires-Dist: matplotlib; extra == "benchmark"
+Requires-Dist: pandas; extra == "benchmark"
+Requires-Dist: folktables; extra == "benchmark"
+Requires-Dist: pot; extra == "benchmark"
+Requires-Dist: hydra; extra == "benchmark"
+Requires-Dist: omegaconf; extra == "benchmark"
 Provides-Extra: ghost
 Requires-Dist: qpsolvers; extra == "ghost"
 Requires-Dist: scipy; extra == "ghost"
 Dynamic: license-file
 
-# Benchmarking Stochastic Approximation Algorithms for Fairness-Constrained Training of Deep Neural Networks
+# humancompatible-train: a package for constrained machine learning
 
 [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) [![Setup](https://github.com/humancompatible/train/actions/workflows/setup.yml/badge.svg)](https://github.com/humancompatible/train/actions/workflows/setup.yml)
 
-The toolkit implements algorithms for constrained training of neural networks based on PyTorch, and inspired by PyTorch's API, as well as a tool to compare stochastic-constrained stochastic optimization algorithms on a _fair learning_ task in the `experiments` folder.
+The toolkit implements algorithms for constrained training of neural networks based on PyTorch, and inspired by PyTorch's API.
+<!-- , as well as a tool to compare stochastic-constrained stochastic optimization algorithms on a _fair learning_ task in the `experiments` folder. -->
 
 ## Table of Contents
 1. [Basic installation instructions](#basic-installation-instructions)
 2. [Using the toolkit](#using-the-toolkit)
-3. [Reproducing the Benchmark](#reproducing-the-benchmark)
-4. [Extending the toolkit](#extending-the-toolkit) <!-- 6. [Citing humancompatible/train](#Citing-humancompatible/train) -->
+3. [Extending the toolkit](#extending-the-toolkit) 
+4. [Reproducing the Benchmark](#reproducing-the-benchmark)
 5. [License and terms of use](#license-and-terms-of-use)
 6. [References](#references)
 
-Humancompatible/train is still under active development! If you find bugs or have feature
+humancompatible-train is still under active development! If you find bugs or have feature
 requests, please file a
 [Github issue](https://github.com/humancompatible/train/issues). 
 
-## Basic installation instructions
+## Installation
+
+Use
+
+```
+pip install humancompatible-train
+```
+
+The only dependencies of this package are `numpy` and `torch`.
+
+## Using the toolkit
+
+The toolkit implements algorithms for constrained training of neural networks based on PyTorch.
+
+The algorithms follow the `dual_step()` - `step()` framework: taking inspiration from PyTorch, the `double_step` does updates related to the dual parameters and prepares for the primal update (by, e.g., saving constraint gradients), and `step()` updates the primal parameters.
+
+In general, your code using `humancompatible-train` would look something like this:
+
+```
+for inputs, labels in dataloader:
+  # inference
+  outputs = model(inputs)
+  # calculate constraints and grads
+  for constraint in constraints:
+      c_eval = constraint(outputs, labels)
+      c_eval.backwards(retain_grad=True)
+      # depending on optimizer, update dual parameters / save constraint gradient / both
+      optimizer.dual_step(c_eval)
+      optimizer.zero_grad()
+  # calculate objective
+  loss = criterion(outputs,labels)
+  loss.backwards()
+  optimizer.step()
+  optimizer.zero_grad()
+```
+
+Our idea is to
+1. Deviate minimally from the usual PyTorch workflow
+2. Make different stochastic-constrained stochastic optimization algorithms nearly interchangable in the code.
+
+### Code examples
+
+You are invited to check out our new API presented in notebooks in the `examples` folder.
+
+*The legacy API used for the benchmark is presented in `examples/_old_/algorithm_demo.ipynb` and `examples/_old_/constraint_demo.ipynb`.*
+
+## Extending the toolkit
+
+### Adding new code
+
+**To add a new algorithm**, you can subclass the PyTorch ```Optimizer``` class and proceed following the API guideline presented above.
+
+## Reproducing the Benchmark
+
+The code used in [our benchmark paper](https://arxiv.org/abs/2507.04033) is not migrated to the new API yet (WIP).
+
+### Basic installation instructions
 The code requires Python version ```3.11```.
 
 1. Create a virtual environment
@@ -72,22 +137,6 @@ after installing requirements.txt; otherwise, the algorithm will run slower. How
 <!-- pip install folktables -->
 <!-- ``` -->
 
-## Using the toolkit
-
-The toolkit implements algorithms for constrained training of neural networks based on PyTorch, and inspired by PyTorch's API.
-
-### Code examples
-
-You are invited to check out the new API presented in notebooks in the `examples` folder.
-
-The algorithms follow the `dual_step()` - `step()` framework: taking inspiration from PyTorch, the `double_step` does updates related to the dual parameters and prepares for the primal update (by, e.g., saving constraint gradients), and `step()` updates the primal parameters.
-
-The idea is to make different algorithms nearly interchangable in the code.
-
-The legacy API used for the benchmark is presented in `examples/_old_/algorithm_demo.ipynb` and `examples/_old_/constraint_demo.ipynb`.
-
-## Reproducing the Benchmark
-
 ### Running the algorithms
 
 The benchmark comprises the following algorithms:
@@ -122,15 +171,10 @@ This repository uses [Hydra](https://hydra.cc/) to manage parameters; see `exper
 ### Producing plots
 The plots and tables like the ones in the paper can be produced using the two notebooks. `experiments/algo_plots.ipynb` houses the convergence plots, and `experiments/model_plots.ipynb` - all the others.
 
-## Extending the toolkit
-
-### Adding new code
-
-**To add a new algorithm**, you can subclass the PyTorch ```Optimizer``` class and proceed following the API guideline presented above.
 
 ## License and terms of use
 
-humancompatible/train is provided under the Apache 2.0 Licence.
+humancompatible-train is provided under the Apache 2.0 Licence.
 
 The benchmark part of the package relies on the Folktables package, provided under MIT Licence.
 It provides code to download data from the American Community Survey
@@ -154,7 +198,7 @@ For more information, see https://www.census.gov/data/developers/about/terms-of-
 
 ## Future work
 
-- Add more algorithms with PyTorch-like API
+- Add more algorithms
 - Add more examples from different fields where constrained training of DNNs is employed
 - Migrate the benchmark to the new API
 
@@ -185,4 +229,3 @@ Huang, Zhang & Alacaoglu (2025) Stochastic Smoothed Primal-Dual Algorithms for N
 
 <a id="4">[4]</a> 
 Huang & Lin (2023) Oracle Complexity of Single-Loop Switching Subgradient Methods for Non-Smooth Weakly Convex Functional Constrained Optimization, Curran Associates Inc..
-

humancompatible_train-0.1.0/PKG-INFO → humancompatible_train-0.1.2/README.md

@@ -1,39 +1,79 @@
-Metadata-Version: 2.4
-Name: humancompatible-train
-Version: 0.1.0
-Summary: PyTorch-based package for constrained training of neural networks
-Author: Gilles Bareilles, Jana Lepsova, Jakub Marecek
-Author-email: Andrii Kliachkin <kliachkin.andrii@gmail.com>
-Maintainer-email: Andrii Kliachkin <kliachkin.andrii@gmail.com>
-Requires-Python: >=3.11
-Description-Content-Type: text/markdown
-License-File: LICENCE.txt
-Requires-Dist: torch
-Requires-Dist: numpy
-Provides-Extra: ghost
-Requires-Dist: qpsolvers; extra == "ghost"
-Requires-Dist: scipy; extra == "ghost"
-Dynamic: license-file
-
-# Benchmarking Stochastic Approximation Algorithms for Fairness-Constrained Training of Deep Neural Networks
+# humancompatible-train: a package for constrained machine learning
 
 [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) [![Setup](https://github.com/humancompatible/train/actions/workflows/setup.yml/badge.svg)](https://github.com/humancompatible/train/actions/workflows/setup.yml)
 
-The toolkit implements algorithms for constrained training of neural networks based on PyTorch, and inspired by PyTorch's API, as well as a tool to compare stochastic-constrained stochastic optimization algorithms on a _fair learning_ task in the `experiments` folder.
+The toolkit implements algorithms for constrained training of neural networks based on PyTorch, and inspired by PyTorch's API.
+<!-- , as well as a tool to compare stochastic-constrained stochastic optimization algorithms on a _fair learning_ task in the `experiments` folder. -->
 
 ## Table of Contents
 1. [Basic installation instructions](#basic-installation-instructions)
 2. [Using the toolkit](#using-the-toolkit)
-3. [Reproducing the Benchmark](#reproducing-the-benchmark)
-4. [Extending the toolkit](#extending-the-toolkit) <!-- 6. [Citing humancompatible/train](#Citing-humancompatible/train) -->
+3. [Extending the toolkit](#extending-the-toolkit) 
+4. [Reproducing the Benchmark](#reproducing-the-benchmark)
 5. [License and terms of use](#license-and-terms-of-use)
 6. [References](#references)
 
-Humancompatible/train is still under active development! If you find bugs or have feature
+humancompatible-train is still under active development! If you find bugs or have feature
 requests, please file a
 [Github issue](https://github.com/humancompatible/train/issues). 
 
-## Basic installation instructions
+## Installation
+
+Use
+
+```
+pip install humancompatible-train
+```
+
+The only dependencies of this package are `numpy` and `torch`.
+
+## Using the toolkit
+
+The toolkit implements algorithms for constrained training of neural networks based on PyTorch.
+
+The algorithms follow the `dual_step()` - `step()` framework: taking inspiration from PyTorch, the `double_step` does updates related to the dual parameters and prepares for the primal update (by, e.g., saving constraint gradients), and `step()` updates the primal parameters.
+
+In general, your code using `humancompatible-train` would look something like this:
+
+```
+for inputs, labels in dataloader:
+  # inference
+  outputs = model(inputs)
+  # calculate constraints and grads
+  for constraint in constraints:
+      c_eval = constraint(outputs, labels)
+      c_eval.backwards(retain_grad=True)
+      # depending on optimizer, update dual parameters / save constraint gradient / both
+      optimizer.dual_step(c_eval)
+      optimizer.zero_grad()
+  # calculate objective
+  loss = criterion(outputs,labels)
+  loss.backwards()
+  optimizer.step()
+  optimizer.zero_grad()
+```
+
+Our idea is to
+1. Deviate minimally from the usual PyTorch workflow
+2. Make different stochastic-constrained stochastic optimization algorithms nearly interchangable in the code.
+
+### Code examples
+
+You are invited to check out our new API presented in notebooks in the `examples` folder.
+
+*The legacy API used for the benchmark is presented in `examples/_old_/algorithm_demo.ipynb` and `examples/_old_/constraint_demo.ipynb`.*
+
+## Extending the toolkit
+
+### Adding new code
+
+**To add a new algorithm**, you can subclass the PyTorch ```Optimizer``` class and proceed following the API guideline presented above.
+
+## Reproducing the Benchmark
+
+The code used in [our benchmark paper](https://arxiv.org/abs/2507.04033) is not migrated to the new API yet (WIP).
+
+### Basic installation instructions
 The code requires Python version ```3.11```.
 
 1. Create a virtual environment
@@ -72,22 +112,6 @@ after installing requirements.txt; otherwise, the algorithm will run slower. How
 <!-- pip install folktables -->
 <!-- ``` -->
 
-## Using the toolkit
-
-The toolkit implements algorithms for constrained training of neural networks based on PyTorch, and inspired by PyTorch's API.
-
-### Code examples
-
-You are invited to check out the new API presented in notebooks in the `examples` folder.
-
-The algorithms follow the `dual_step()` - `step()` framework: taking inspiration from PyTorch, the `double_step` does updates related to the dual parameters and prepares for the primal update (by, e.g., saving constraint gradients), and `step()` updates the primal parameters.
-
-The idea is to make different algorithms nearly interchangable in the code.
-
-The legacy API used for the benchmark is presented in `examples/_old_/algorithm_demo.ipynb` and `examples/_old_/constraint_demo.ipynb`.
-
-## Reproducing the Benchmark
-
 ### Running the algorithms
 
 The benchmark comprises the following algorithms:
@@ -122,15 +146,10 @@ This repository uses [Hydra](https://hydra.cc/) to manage parameters; see `exper
 ### Producing plots
 The plots and tables like the ones in the paper can be produced using the two notebooks. `experiments/algo_plots.ipynb` houses the convergence plots, and `experiments/model_plots.ipynb` - all the others.
 
-## Extending the toolkit
-
-### Adding new code
-
-**To add a new algorithm**, you can subclass the PyTorch ```Optimizer``` class and proceed following the API guideline presented above.
 
 ## License and terms of use
 
-humancompatible/train is provided under the Apache 2.0 Licence.
+humancompatible-train is provided under the Apache 2.0 Licence.
 
 The benchmark part of the package relies on the Folktables package, provided under MIT Licence.
 It provides code to download data from the American Community Survey
@@ -154,7 +173,7 @@ For more information, see https://www.census.gov/data/developers/about/terms-of-
 
 ## Future work
 
-- Add more algorithms with PyTorch-like API
+- Add more algorithms
 - Add more examples from different fields where constrained training of DNNs is employed
 - Migrate the benchmark to the new API
 
@@ -184,5 +203,4 @@ Facchinei & Kungurtsev (2023) Stochastic Approximation for Expectation Objective
 Huang, Zhang & Alacaoglu (2025) Stochastic Smoothed Primal-Dual Algorithms for Nonconvex Optimization with Linear Inequality Constraints, arXiv.
 
 <a id="4">[4]</a> 
-Huang & Lin (2023) Oracle Complexity of Single-Loop Switching Subgradient Methods for Non-Smooth Weakly Convex Functional Constrained Optimization, Curran Associates Inc..
-
+Huang & Lin (2023) Oracle Complexity of Single-Loop Switching Subgradient Methods for Non-Smooth Weakly Convex Functional Constrained Optimization, Curran Associates Inc..

{humancompatible_train-0.1.0 → humancompatible_train-0.1.2}/experiments/run_folktables.py

@@ -12,7 +12,7 @@ from omegaconf import DictConfig, OmegaConf
 from torch import nn, tensor
 from utils.load_folktables import prepare_folktables_multattr
 from utils.network import SimpleNet
-from humancompatible.train.algorithms.utils import net_grads_to_tensor
+from humancompatible.train.benchmark.algorithms.utils import net_grads_to_tensor
 from itertools import combinations
 from humancompatible.train.fairness.constraints import FairnessConstraint
 
@@ -177,7 +177,7 @@ def run(cfg: DictConfig) -> None:
         net = SimpleNet(in_shape=X_test.shape[1], out_shape=1, dtype=DTYPE).to(device)
 
         optimizer_name = cfg.alg.import_name
-        module = importlib.import_module("humancompatible.train.algorithms")
+        module = importlib.import_module("humancompatible.train.benchmark.algorithms")
         Optimizer = getattr(module, optimizer_name)
 
         optimizer = Optimizer(net, train_ds, loss_fn, c)

humancompatible_train-0.1.2/humancompatible/train/algorithms/test/test_ssl_alm.py

@@ -0,0 +1,144 @@
+import unittest
+import torch
+from torch import Tensor
+from humancompatible.train.algorithms import SSLALM
+
+class TestSSLALM(unittest.TestCase):
+    def setUp(self):
+        # Simple model for testing
+        self.model = torch.nn.Linear(2, 1)
+        self.params = list(self.model.parameters())
+        self.m = 2  # Number of constraints
+        self.optimizer = SSLALM(
+            self.params,
+            m=self.m,
+            lr=0.01,
+            dual_lr=0.01,
+            dual_bound=100.0,
+            rho=1.0,
+            mu=2.0,
+            beta=0.5,
+        )
+
+    def test_initialization(self):
+        # Test if the optimizer is initialized correctly
+        self.assertEqual(len(self.optimizer.param_groups), 1)
+        self.assertEqual(self.optimizer.m, self.m)
+        self.assertEqual(self.optimizer.dual_lr, 0.01)
+        self.assertEqual(self.optimizer.dual_bound, 100.0)
+        self.assertEqual(self.optimizer.rho, 1.0)
+        self.assertEqual(self.optimizer.mu, 2.0)
+        self.assertEqual(self.optimizer.beta, 0.5)
+        self.assertTrue(isinstance(self.optimizer._dual_vars, Tensor))
+        self.assertEqual(self.optimizer._dual_vars.shape, (self.m,))
+
+    def test_dual_step(self):
+        # Test dual variable update
+        c_val = torch.tensor([0.5, 0.1])
+        self.optimizer.dual_step(0, c_val[0])
+        self.assertEqual(self.optimizer._dual_vars[0], 0.005)  # 0 + 0.01 * 0.5
+        self.optimizer.dual_step(1, c_val[1])
+        self.assertEqual(self.optimizer._dual_vars[1], 0.001)  # 0 + 0.01 * 0.1
+
+    def test_dual_bound(self):
+        # Test dual variable bounding
+        self.optimizer._dual_vars = torch.tensor([101.0, -1.0])
+        c_val = torch.tensor([1.0, -1.0])
+        self.optimizer.dual_step(0, c_val[0])
+        self.optimizer.dual_step(1, c_val[1])
+        self.assertEqual(self.optimizer._dual_vars[0], 0.0)  # Should be zeroed out
+        self.assertEqual(self.optimizer._dual_vars[1], 0.0)  # Should be zeroed out
+
+# ADD TEST DEALING WITH CONSTRAINTS THAT DONT USE SOME OF THE PARAMS
+
+    def test_step(self):
+        # Test primal parameter update
+        # Mock gradients and constraint gradients
+        p_pre_step = {}
+        for p in self.params:
+            p.grad = torch.ones_like(p)
+            p_pre_step[p] = p.detach().clone()
+
+        c_val = torch.tensor([0.1, -0.1])
+        c_grads = {p: [torch.ones_like(p) for _ in c_val] for p in self.params}
+        
+        for p in self.params:
+            self.optimizer.state[p]["c_grad"] = [g.clone() for g in c_grads[p]]
+            self.optimizer.state[p]["smoothing"] = p.detach().clone()
+        
+        self.optimizer._dual_vars = torch.ones(2)
+        G = self.optimizer.step(c_val)
+        # Check if G is computed and parameters are updated
+        self.assertEqual(len(G), len(self.params))
+        for i, p in enumerate(self.params):
+            self.assertTrue(
+                torch.equal(
+                    G[i],
+                    (
+                        p.grad +
+                        sum(_lambda  * c_grads[p][j] for j, _lambda in enumerate(self.optimizer._dual_vars)) +
+                        sum([c_grads[p][j] * cv for j, cv in enumerate(c_val)])
+                    )
+                )
+            )
+
+    def test_step_with_none_c_val(self):
+        # Test step with None c_val (should use self.c_vals)
+        c_val = torch.tensor([0.1, -0.1])
+        self.optimizer.dual_step(0, c_val[0])
+        self.optimizer.dual_step(1, c_val[1])
+
+        p_pre_step = {}
+        for p in self.params:
+            p.grad = torch.ones_like(p)
+            p_pre_step[p] = p.detach().clone()
+
+        c_grads = {p: [torch.ones_like(p) for _ in c_val] for p in self.params}
+        for p in self.params:
+            self.optimizer.state[p]["c_grad"] = [g.clone() for g in c_grads[p]]
+            self.optimizer.state[p]["smoothing"] = p.detach().clone()
+        
+        G = self.optimizer.step()
+        # Check if G is computed and parameters are updated
+        self.assertEqual(len(G), len(self.params))
+        for i, p in enumerate(self.params):
+            self.assertTrue(
+                torch.equal(
+                    G[i],
+                    (
+                        p.grad +
+                        sum(_lambda  * c_grads[p][j] for j, _lambda in enumerate(self.optimizer._dual_vars)) +
+                        sum([c_grads[p][j] * cv for j, cv in enumerate(c_val)])
+                    )
+                )
+            )
+
+    def test_step_with_invalid_c_val(self):
+        # Test step with invalid c_val (wrong shape)
+        with self.assertRaises(ValueError):
+            self.optimizer.step(torch.tensor([0.1]))
+
+    def test_smoothing_update(self):
+        # Test smoothing term update
+        p_pre_step = {}
+        for p in self.params:
+            p.grad = torch.ones_like(p)
+            p_pre_step[p] = p.detach().clone()
+        c_val = torch.tensor([0.1, -0.1])
+        self.optimizer.step(c_val)
+        for p in self.params:
+            state = self.optimizer.state[p]
+            self.assertTrue("smoothing" in state)
+            self.assertTrue(torch.all(state["smoothing"] == p_pre_step[p]))
+
+    def test_error_handling(self):
+        # Test error handling for invalid inputs
+        with self.assertRaises(ValueError):
+            SSLALM(self.params, m=self.m, lr=-0.01)
+        with self.assertRaises(ValueError):
+            SSLALM(self.params, m=self.m, dual_lr=-0.01)
+        with self.assertRaises(ValueError):
+            SSLALM(self.params, m=self.m, init_dual_vars=torch.tensor([1.0]))
+
+if __name__ == "__main__":
+    unittest.main()

humancompatible_train-0.1.2/humancompatible/train/algorithms/test/test_ssw.py

@@ -0,0 +1,64 @@
+import unittest
+import torch
+from humancompatible.train.algorithms import SSG
+
+class TestSSG(unittest.TestCase):
+    def setUp(self):
+        # Simple model for testing
+        self.model = torch.nn.Linear(2, 1)
+        self.params = list(self.model.parameters())
+        self.m = 1  # Number of constraints
+        self.optimizer = SSG(
+            self.params,
+            m=self.m,
+            lr=0.01,
+            dual_lr=0.01
+        )
+
+    def test_initialization(self):
+        # Test if the optimizer is initialized correctly
+        self.assertEqual(len(self.optimizer.param_groups), 1)
+        self.assertEqual(self.optimizer.m, self.m)
+        self.assertEqual(self.optimizer.dual_lr, 0.01)
+
+    def test_dual_step(self):
+        # Test dual step saving constraint gradients
+        for p in self.params:
+            p.grad = torch.ones_like(p)
+        self.optimizer.dual_step(0)
+        
+        for p in self.params:
+            state = self.optimizer.state[p]
+            self.assertTrue(state["c_grad"] is not None)
+
+    def test_step_obj(self):
+        p_pre_step = {}
+        for p in self.params:
+            p.grad = torch.ones_like(p)
+            self.optimizer.state[p]['c_grad'] = [-1.*torch.ones_like(p)]
+            p_pre_step[p] = p.detach().clone()
+
+        self.optimizer.step(c_val=torch.tensor(-1.))
+
+        for p in self.params:
+            self.assertTrue(
+                torch.all(
+                        p == p_pre_step[p]-0.01*torch.ones_like(p)
+                    )
+                )
+
+    def test_step_c(self):
+        p_pre_step = {}
+        for p in self.params:
+            p.grad = torch.ones_like(p)
+            self.optimizer.state[p]['c_grad'] = [-1.*torch.ones_like(p)]
+            p_pre_step[p] = p.detach().clone()
+
+        self.optimizer.step(c_val=torch.tensor(1.))
+
+        for p in self.params:
+            self.assertTrue(
+                torch.all(
+                        p == p_pre_step[p]+0.01*torch.ones_like(p)
+                    )
+                )

humancompatible_train-0.1.2/humancompatible/train/fairness/utils/tests/test_balanced_batch_sampler.py

@@ -0,0 +1,78 @@
+import unittest
+import torch
+from torch.utils.data import TensorDataset, Subset, DataLoader
+from humancompatible.train.fairness.utils import BalancedBatchSampler
+
+class TestBalancedBatchSampler(unittest.TestCase):
+    def setUp(self):
+        self.data = torch.tensor([[i, i+1] for i in range(10)])
+        self.labels = torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 2, 2])
+        self.dataset = TensorDataset(self.data, self.labels)
+        self.subset_indices = [
+            [0, 1],  # Class 0
+            [2, 3, 4],  # Class 1
+            [5, 6, 7, 8, 9],  # Class 2
+        ]
+        self.subset_onehot = torch.tensor([
+            [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
+            [0, 0, 1, 1, 1, 0, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
+        ]).T
+
+    def test_batch_size_divisible(self):
+        with self.assertRaises(AssertionError):
+            BalancedBatchSampler(subgroup_indices=self.subset_indices, batch_size=4, drop_last=True)
+
+    def test_onehot_init(self):
+        sampler = BalancedBatchSampler(subgroup_onehot=self.subset_onehot, batch_size=3)
+        self.assertListEqual(
+            [i.tolist() for i in sampler.subset_indices],
+            self.subset_indices
+        )
+
+    def test_iter(self):
+        sampler = BalancedBatchSampler(subgroup_indices=self.subset_indices, batch_size=6, drop_last=True)
+        batches = list(sampler)
+        self.assertEqual(len(batches), 1)  # Only 1 full batch of size 6 (2+2+2)
+        self.assertEqual(len(batches[0]), 6)
+
+    def test_len_drop_last_true(self):
+        sampler = BalancedBatchSampler(subgroup_indices=self.subset_indices, batch_size=6, drop_last=True)
+        self.assertEqual(len(sampler), 1)
+
+    def test_balanced_batches(self):
+        sampler = BalancedBatchSampler(subgroup_indices=self.subset_indices, batch_size=6, drop_last=True)
+        batch = next(iter(sampler))
+        # Check that each subset contributes 2 samples
+        self.assertEqual(len([i for i in batch if i in self.subset_indices[0]]), 2)
+        self.assertEqual(len([i for i in batch if i in self.subset_indices[1]]), 2)
+        self.assertEqual(len([i for i in batch if i in self.subset_indices[2]]), 2)
+
+class TestDataLoaderIntegration(unittest.TestCase):
+    def setUp(self):
+        self.data = torch.tensor([[i, i+1] for i in range(10)])
+        self.labels = torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 2, 2])
+        self.dataset = TensorDataset(self.data, self.labels)
+        self.subset_indices = [
+            [0, 1],  # Class 0
+            [2, 3, 4],  # Class 1
+            [5, 6, 7, 8, 9],  # Class 2
+        ]
+        self.subsets = [Subset(self.dataset, indices) for indices in self.subset_indices]
+
+    def test_dataloader(self):
+        sampler = BalancedBatchSampler(subgroup_indices=self.subset_indices, batch_size=6, drop_last=True)
+        dataloader = DataLoader(
+            self.dataset,
+            batch_sampler=sampler
+        )
+        batch_data, batch_labels = next(iter(dataloader))
+        self.assertEqual(batch_data.shape, (6, 2))
+        self.assertEqual(len(batch_labels), 6)
+        # Check balance: 2 samples from each class
+        self.assertEqual((batch_labels == 0).sum().item(), 2)
+        self.assertEqual((batch_labels == 1).sum().item(), 2)
+        self.assertEqual((batch_labels == 2).sum().item(), 2)
+
+if __name__ == "__main__":
+    unittest.main()