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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: humancompatible-train
-Version: 0.1.2
+Version: 0.1.3
 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,6 +10,13 @@ Description-Content-Type: text/markdown
 License-File: LICENCE.txt
 Requires-Dist: torch
 Requires-Dist: numpy
+Provides-Extra: examples
+Requires-Dist: ipykernel; extra == "examples"
+Requires-Dist: ipympl; extra == "examples"
+Requires-Dist: fairret; extra == "examples"
+Requires-Dist: folktables; extra == "examples"
+Requires-Dist: scikit-learn; extra == "examples"
+Requires-Dist: matplotlib; extra == "examples"
 Provides-Extra: benchmark
 Requires-Dist: fairret; extra == "benchmark"
 Requires-Dist: matplotlib; extra == "benchmark"
@@ -31,22 +38,23 @@ The toolkit implements algorithms for constrained training of neural networks ba
 <!-- , 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. [Extending the toolkit](#extending-the-toolkit) 
+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
 requests, please file a
-[Github issue](https://github.com/humancompatible/train/issues). 
+[Github issue](https://github.com/humancompatible/train/issues).
 
 ## Installation
 
 Use
 
-```
+```bash
 pip install humancompatible-train
 ```
 
@@ -56,11 +64,11 @@ The only dependencies of this package are `numpy` and `torch`.
 
 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.
+The algorithms follow the `dual_step()` - `step()` framework: taking inspiration from PyTorch, the `dual_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:
 
-```
+```python
 for inputs, labels in dataloader:
   # inference
   outputs = model(inputs)
@@ -79,6 +87,7 @@ for inputs, labels in dataloader:
 ```
 
 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.
 
@@ -86,6 +95,12 @@ Our idea is to
 
 You are invited to check out our new API presented in notebooks in the `examples` folder.
 
+The example notebooks have additional dependencies, such as `fairret`. To install those, run
+
+```
+pip install humancompatible-train[examples]
+```
+
 *The legacy API used for the benchmark is presented in `examples/_old_/algorithm_demo.ipynb` and `examples/_old_/constraint_demo.ipynb`.*
 
 ## Extending the toolkit
@@ -99,21 +114,27 @@ You are invited to check out our new API presented in notebooks in the `examples
 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
 
 **bash** (Linux)
+
 ```
 python3.11 -m venv fairbenchenv
 source fairbenchenv/bin/activate
 ```
+
 **cmd** (Windows)
+
 ```
 python -m venv fairbenchenv
 fairbenchenv\Scripts\activate.bat
 ```
+
 2. Install from source.
+
 ```
 git clone https://github.com/humancompatible/train.git
 cd train
@@ -122,11 +143,12 @@ pip install .
 ```
 
 If you wish to edit the code of the algorithms, install as an editable package:
+
 ```
 pip install -e .
 ```
 
-__Warning__: it is recommended to use Stochastic Ghost with the mkl-accelerated version of the scipy package with Stochastic Ghost; to install it, run
+**Warning**: it is recommended to use Stochastic Ghost with the mkl-accelerated version of the scipy package with Stochastic Ghost; to install it, run
 
 ```pip install --force-reinstall -i https://software.repos.intel.com/python/pypi scipy```
 
@@ -140,12 +162,14 @@ after installing requirements.txt; otherwise, the algorithm will run slower. How
 ### Running the algorithms
 
 The benchmark comprises the following algorithms:
+
 - Stochastic Ghost [[2]](#2),
 - SSL-ALM [[3]](#3),
 - Stochastic Switching Subgradient [[4]](#4).
 
 To reproduce the experiments of the paper, run the following:
-``` python
+
+```
 cd experiments
 python run_folktables.py data=folktables alg=sslalm
 python run_folktables.py data=folktables alg=alm
@@ -154,23 +178,25 @@ python run_folktables.py data=folktables alg=ssg
 python run_folktables.py data=folktables alg=sgd     # baseline, no fairness
 python run_folktables.py data=folktables alg=fairret # baseline, fairness with regularizer
 ```
+
 Each command will start 10 runs of the `alg`, 30 seconds each.
 The results will be saved to `experiments/utils/saved_models` and `experiments/utils/exp_results`.
 <!-- In the repository, we include the configuration needed to reproduce the experiments in the paper. To do so, go to `experiments` and run `python run_folktables.py data=folktables alg=sslalm`. -->
 <!-- Repeat for the other algorithms by changing the `alg` parameter. -->
 
-This repository uses [Hydra](https://hydra.cc/) to manage parameters; see `experiments/conf` for configuration files. 
-* To change the parameters of the experiment, such as the number of runs for each algorithm, run time, the dataset used (*note: for now supports only Folktables*) - use `experiment.yaml`. 
-* To change the dataset settings - such as file location - or do dataset-specific adjustments - such as the configuration of the protected attributes - use `data/{dataset_name}.yaml`
-* To change algorithm hyperparameters, use `alg/{algorithm_name}.yaml`.
-* To change constraint hyperparameters, use `constraint/{constraint_name}.yaml`
+This repository uses [Hydra](https://hydra.cc/) to manage parameters; see `experiments/conf` for configuration files.
+
+- To change the parameters of the experiment, such as the number of runs for each algorithm, run time, the dataset used (*note: for now supports only Folktables*) - use `experiment.yaml`.
+- To change the dataset settings - such as file location - or do dataset-specific adjustments - such as the configuration of the protected attributes - use `data/{dataset_name}.yaml`
+- To change algorithm hyperparameters, use `alg/{algorithm_name}.yaml`.
+- To change constraint hyperparameters, use `constraint/{constraint_name}.yaml`
 
 <!-- ; it is installed as one of the dependencies. -->
 <!-- To learn more about using Hydra, please check out the [official tutorial](https://hydra.cc/docs/tutorials/basic/your_first_app). -->
 
 ### 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.
 
+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.
 
 ## License and terms of use
 
@@ -180,7 +206,7 @@ The benchmark part of the package relies on the Folktables package, provided und
 It provides code to download data from the American Community Survey
 (ACS) Public Use Microdata Sample (PUMS) files managed by the US Census Bureau.
 The data itself is governed by the terms of use provided by the Census Bureau.
-For more information, see https://www.census.gov/data/developers/about/terms-of-service.html
+For more information, see <https://www.census.gov/data/developers/about/terms-of-service.html>
 
 <!-- ## Cite this work -->
 
@@ -206,7 +232,7 @@ For more information, see https://www.census.gov/data/developers/about/terms-of-
 
 If you use this work, we encourage you to cite [our paper](https://arxiv.org/abs/2507.04033),
 
-```
+```bibtex
 @misc{kliachkin2025benchmarkingstochasticapproximationalgorithms,
       title={Benchmarking Stochastic Approximation Algorithms for Fairness-Constrained Training of Deep Neural Networks}, 
       author={Andrii Kliachkin and Jana Lepšová and Gilles Bareilles and Jakub Mareček},
@@ -218,14 +244,14 @@ If you use this work, we encourage you to cite [our paper](https://arxiv.org/abs
 }
 ```
 
-<a id="1">[1]</a> 
+<a id="1">[1]</a>
 Ding, Hardt & Miller et al. (2021) Retiring Adult: New Datasets for Fair Machine Learning, Curran Associates, Inc..
 
-<a id="2">[2]</a> 
+<a id="2">[2]</a>
 Facchinei & Kungurtsev (2023) Stochastic Approximation for Expectation Objective and Expectation Inequality-Constrained Nonconvex Optimization, arXiv.
 
-<a id="3">[3]</a> 
+<a id="3">[3]</a>
 Huang, Zhang & Alacaoglu (2025) Stochastic Smoothed Primal-Dual Algorithms for Nonconvex Optimization with Linear Inequality Constraints, arXiv.
 
-<a id="4">[4]</a> 
+<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.2 → humancompatible_train-0.1.3}/README.md

@@ -6,22 +6,23 @@ The toolkit implements algorithms for constrained training of neural networks ba
 <!-- , 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. [Extending the toolkit](#extending-the-toolkit) 
+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
 requests, please file a
-[Github issue](https://github.com/humancompatible/train/issues). 
+[Github issue](https://github.com/humancompatible/train/issues).
 
 ## Installation
 
 Use
 
-```
+```bash
 pip install humancompatible-train
 ```
 
@@ -31,11 +32,11 @@ The only dependencies of this package are `numpy` and `torch`.
 
 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.
+The algorithms follow the `dual_step()` - `step()` framework: taking inspiration from PyTorch, the `dual_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:
 
-```
+```python
 for inputs, labels in dataloader:
   # inference
   outputs = model(inputs)
@@ -54,6 +55,7 @@ for inputs, labels in dataloader:
 ```
 
 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.
 
@@ -61,6 +63,12 @@ Our idea is to
 
 You are invited to check out our new API presented in notebooks in the `examples` folder.
 
+The example notebooks have additional dependencies, such as `fairret`. To install those, run
+
+```
+pip install humancompatible-train[examples]
+```
+
 *The legacy API used for the benchmark is presented in `examples/_old_/algorithm_demo.ipynb` and `examples/_old_/constraint_demo.ipynb`.*
 
 ## Extending the toolkit
@@ -74,21 +82,27 @@ You are invited to check out our new API presented in notebooks in the `examples
 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
 
 **bash** (Linux)
+
 ```
 python3.11 -m venv fairbenchenv
 source fairbenchenv/bin/activate
 ```
+
 **cmd** (Windows)
+
 ```
 python -m venv fairbenchenv
 fairbenchenv\Scripts\activate.bat
 ```
+
 2. Install from source.
+
 ```
 git clone https://github.com/humancompatible/train.git
 cd train
@@ -97,11 +111,12 @@ pip install .
 ```
 
 If you wish to edit the code of the algorithms, install as an editable package:
+
 ```
 pip install -e .
 ```
 
-__Warning__: it is recommended to use Stochastic Ghost with the mkl-accelerated version of the scipy package with Stochastic Ghost; to install it, run
+**Warning**: it is recommended to use Stochastic Ghost with the mkl-accelerated version of the scipy package with Stochastic Ghost; to install it, run
 
 ```pip install --force-reinstall -i https://software.repos.intel.com/python/pypi scipy```
 
@@ -115,12 +130,14 @@ after installing requirements.txt; otherwise, the algorithm will run slower. How
 ### Running the algorithms
 
 The benchmark comprises the following algorithms:
+
 - Stochastic Ghost [[2]](#2),
 - SSL-ALM [[3]](#3),
 - Stochastic Switching Subgradient [[4]](#4).
 
 To reproduce the experiments of the paper, run the following:
-``` python
+
+```
 cd experiments
 python run_folktables.py data=folktables alg=sslalm
 python run_folktables.py data=folktables alg=alm
@@ -129,23 +146,25 @@ python run_folktables.py data=folktables alg=ssg
 python run_folktables.py data=folktables alg=sgd     # baseline, no fairness
 python run_folktables.py data=folktables alg=fairret # baseline, fairness with regularizer
 ```
+
 Each command will start 10 runs of the `alg`, 30 seconds each.
 The results will be saved to `experiments/utils/saved_models` and `experiments/utils/exp_results`.
 <!-- In the repository, we include the configuration needed to reproduce the experiments in the paper. To do so, go to `experiments` and run `python run_folktables.py data=folktables alg=sslalm`. -->
 <!-- Repeat for the other algorithms by changing the `alg` parameter. -->
 
-This repository uses [Hydra](https://hydra.cc/) to manage parameters; see `experiments/conf` for configuration files. 
-* To change the parameters of the experiment, such as the number of runs for each algorithm, run time, the dataset used (*note: for now supports only Folktables*) - use `experiment.yaml`. 
-* To change the dataset settings - such as file location - or do dataset-specific adjustments - such as the configuration of the protected attributes - use `data/{dataset_name}.yaml`
-* To change algorithm hyperparameters, use `alg/{algorithm_name}.yaml`.
-* To change constraint hyperparameters, use `constraint/{constraint_name}.yaml`
+This repository uses [Hydra](https://hydra.cc/) to manage parameters; see `experiments/conf` for configuration files.
+
+- To change the parameters of the experiment, such as the number of runs for each algorithm, run time, the dataset used (*note: for now supports only Folktables*) - use `experiment.yaml`.
+- To change the dataset settings - such as file location - or do dataset-specific adjustments - such as the configuration of the protected attributes - use `data/{dataset_name}.yaml`
+- To change algorithm hyperparameters, use `alg/{algorithm_name}.yaml`.
+- To change constraint hyperparameters, use `constraint/{constraint_name}.yaml`
 
 <!-- ; it is installed as one of the dependencies. -->
 <!-- To learn more about using Hydra, please check out the [official tutorial](https://hydra.cc/docs/tutorials/basic/your_first_app). -->
 
 ### 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.
 
+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.
 
 ## License and terms of use
 
@@ -155,7 +174,7 @@ The benchmark part of the package relies on the Folktables package, provided und
 It provides code to download data from the American Community Survey
 (ACS) Public Use Microdata Sample (PUMS) files managed by the US Census Bureau.
 The data itself is governed by the terms of use provided by the Census Bureau.
-For more information, see https://www.census.gov/data/developers/about/terms-of-service.html
+For more information, see <https://www.census.gov/data/developers/about/terms-of-service.html>
 
 <!-- ## Cite this work -->
 
@@ -181,7 +200,7 @@ For more information, see https://www.census.gov/data/developers/about/terms-of-
 
 If you use this work, we encourage you to cite [our paper](https://arxiv.org/abs/2507.04033),
 
-```
+```bibtex
 @misc{kliachkin2025benchmarkingstochasticapproximationalgorithms,
       title={Benchmarking Stochastic Approximation Algorithms for Fairness-Constrained Training of Deep Neural Networks}, 
       author={Andrii Kliachkin and Jana Lepšová and Gilles Bareilles and Jakub Mareček},
@@ -193,14 +212,14 @@ If you use this work, we encourage you to cite [our paper](https://arxiv.org/abs
 }
 ```
 
-<a id="1">[1]</a> 
+<a id="1">[1]</a>
 Ding, Hardt & Miller et al. (2021) Retiring Adult: New Datasets for Fair Machine Learning, Curran Associates, Inc..
 
-<a id="2">[2]</a> 
+<a id="2">[2]</a>
 Facchinei & Kungurtsev (2023) Stochastic Approximation for Expectation Objective and Expectation Inequality-Constrained Nonconvex Optimization, arXiv.
 
-<a id="3">[3]</a> 
+<a id="3">[3]</a>
 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..
+<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.2 → humancompatible_train-0.1.3}/humancompatible/train/algorithms/__init__.py

@@ -1,4 +1,5 @@
 from .ssl_alm import SSLALM
 from .ssw import SSG
+from .ssl_alm_adam import SSLALM_Adam
 
 __all__ = ["SSLALM", "SSG"]

{humancompatible_train-0.1.2 → humancompatible_train-0.1.3}/humancompatible/train/algorithms/ssl_alm.py

@@ -2,7 +2,6 @@ from typing import Iterable, Optional, Union
 
 import torch
 from torch import Tensor
-
 from torch.optim.optimizer import Optimizer, _use_grad_for_differentiable
 
 class SSLALM(Optimizer):

humancompatible_train-0.1.3/humancompatible/train/algorithms/ssl_alm_adam.py

@@ -0,0 +1,330 @@
+from typing import Iterable, Optional, Union
+
+import torch
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, _use_grad_for_differentiable
+
+class SSLALM_Adam(Optimizer):
+    def __init__(
+        self,
+        params,
+        m: int,
+        # tau in paper
+        lr: Union[float, Tensor] = 5e-2,
+        # eta in paper
+        dual_lr: Union[
+            float, Tensor
+        ] = 5e-2,  # keep as tensor for different learning rates for different constraints in the future? idk
+        dual_bound : Union[
+            float, Tensor
+        ] = 100,
+        # penalty term multiplier
+        rho: float = 1.0,
+        # smoothing term multiplier
+        mu: float = 2.0,
+        # smoothing term update multiplier
+        beta: float = 0.5,
+        beta1: float = 0.9,
+        beta2: float = 0.999,
+        eps: float = 1e-8,
+        amsgrad: bool = False,
+        *,
+        init_dual_vars: Optional[Tensor] = None,
+        # whether some of the dual variables should not be updated
+        fix_dual_vars: Optional[Tensor] = None,
+        differentiable: bool = False,
+        # custom_project_fn: Optional[Callable] = project_fn
+    ):
+        if isinstance(lr, torch.Tensor) and lr.numel() != 1:
+            raise ValueError("Tensor lr must be 1-element")
+        if isinstance(dual_lr, torch.Tensor) and lr.numel() != 1:
+            raise ValueError("Tensor dual_lr must be 1-element")
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if dual_lr < 0.0:
+            raise ValueError(f"Invalid dual learning rate: {dual_lr}")
+        if init_dual_vars is not None and len(init_dual_vars) != m:
+            raise ValueError(
+                f"init_dual_vars should be of length m: expected {m}, got {len(init_dual_vars)}"
+            )
+        if fix_dual_vars is not None:
+            raise NotImplementedError()
+        if init_dual_vars is None and fix_dual_vars is not None:
+            raise ValueError(
+                f"if fix_dual_vars is not None, init_dual_vars should not be None."
+            )
+
+        if differentiable:
+            raise NotImplementedError("TorchSSLALM does not support differentiable")
+
+        defaults = dict(
+            lr=lr,
+            dual_lr=dual_lr,
+            rho=rho,
+            mu=mu,
+            beta=beta,
+            amsgrad=amsgrad,
+            differentiable=differentiable,
+            # custom_project_fn=custom_project_fn
+        )
+
+        super().__init__(params, defaults)
+
+        # self.param_groups.append()
+
+        self.m = m
+        self.dual_lr = dual_lr
+        self.dual_bound = dual_bound
+        self.rho = rho
+        self.beta = beta
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.mu = mu
+        self.c_vals: list[Union[float, Tensor]] = []
+        self._c_val_average = [None]
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.eps = eps
+        # essentially, move everything here to self.state[param_group]
+        # self.state[param_group]['smoothing_avg'] <= z for that param_group;
+        # ...['grad'] <= grad w.r.t. that param_group
+        # ...['G'] <= G w.r.t. that param_group // idk if necessary
+        # ...['c_grad'][c_i] <= grad of ith constraint w.r.t. that group<w
+        if init_dual_vars is not None:
+            self._dual_vars = init_dual_vars
+        else:
+            self._dual_vars = torch.zeros(m, requires_grad=False)
+
+    def _init_group(
+        self,
+        group,
+        params,
+        grads,
+        c_grads,
+        exp_avgs,
+        exp_avg_sqs,
+        max_exp_avg_sqs,
+        state_steps,
+        smoothing
+    ):
+        # SHOULDN'T calculate values, only set them from the state of the respective param_group
+        # calculations only happen in step() (or rather in the func version of step)
+        has_sparse_grad = False
+
+        for p in group["params"]:
+            state = self.state[p]
+
+            params.append(p)
+            
+            # load z (smoothing term)
+            # Lazy state initialization
+            if len(state) == 0:
+                state["smoothing"] = p.detach().clone()
+                state["c_grad"] = []
+                
+                state["step"] = (
+                    torch.tensor(0.0)
+                )
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                
+                if group["amsgrad"]:
+                    # raise NotImplementedError()
+                    # Maintains max of all exp. moving avg. of sq. grad. values
+                    state["max_exp_avg_sq"] = torch.zeros_like(
+                        p, memory_format=torch.preserve_format
+                    )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+
+            if group["amsgrad"]:
+                max_exp_avg_sqs.append(state["max_exp_avg_sq"])
+            if group["differentiable"] and state["step"].requires_grad:
+                raise RuntimeError(
+                    "`requires_grad` is not supported for `step` in differentiable mode"
+                )
+
+            smoothing.append(state.get("smoothing"))
+            
+            state_steps.append(state["step"])
+
+            grads.append(p.grad)
+            c_grads.append(state.get("c_grad"))
+        return has_sparse_grad
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("amsgrad", False)
+            group.setdefault("maximize", False)
+            group.setdefault("foreach", None)
+            group.setdefault("capturable", False)
+            group.setdefault("differentiable", False)
+            group.setdefault("decoupled_weight_decay", False)
+
+    def dual_step(self, i: int, c_val: Tensor):
+        r"""Perform an update of the dual parameters.
+        Also saves constraint gradient for weight update. To be called BEFORE :func:`step` in an iteration!
+
+        Args:
+            i (int): index of the constraint
+            c_val (Tensor): an estimate of the value of the constraint at which the gradient was computed; used for dual parameter update
+        """
+
+        # c_vals is cleaned in step()
+        self.c_vals.append(c_val.detach())
+
+        # update dual multipliers
+        dual_update_tensor = torch.zeros_like(self._dual_vars)
+        dual_update_tensor[i] = self.dual_lr * c_val
+        self._dual_vars.add_(dual_update_tensor)
+        for i in range(len(self._dual_vars)):
+            if self._dual_vars[i] >= self.dual_bound or self._dual_vars[i] < 0:
+                self._dual_vars[i].zero_()
+
+        # save constraint grad
+        for group in self.param_groups:
+            params: list[Tensor] = []
+            grads: list[Tensor] = []
+            c_grads: list[Tensor] = []
+            smoothing: list[Tensor] = []
+            exp_avgs: list[Tensor] = []
+            exp_avg_sqs: list[Tensor] = []
+            max_exp_avg_sqs: list[Tensor] = []
+            state_steps: list[Tensor] = []
+            _ = self._init_group(
+                group,
+                params,
+                grads,
+                c_grads,
+                exp_avgs,
+                exp_avg_sqs,
+                max_exp_avg_sqs,
+                state_steps,
+                smoothing
+            )
+
+            for p in group["params"]:
+                state = self.state[p]
+                # state['c_grad'] is cleaned in step()
+                # so it is always empty on dual_step()
+                state["c_grad"].append(p.grad)
+
+    @_use_grad_for_differentiable
+    def step(self, c_val: Union[Iterable | Tensor] = None):
+        r"""Perform an update of the primal parameters (network weights & slack variables). To be called AFTER :func:`dual_step` in an iteration!
+
+        Args:
+            c_val (Tensor): an Iterable of estimates of values of **ALL** constraints; used for primal parameter update.
+                Ideally, must be evaluated on an independent sample from the one used in :func:`dual_step`
+        """
+        
+        if c_val is None:
+            c_val = self.c_vals
+        if isinstance(c_val, Iterable) and not isinstance(c_val, torch.Tensor):
+            # if len(c_val) == 1 and isinstance(c_val[0], torch.Tensor):
+            #     c_val = c_val[0]
+            # else:
+            c_val = torch.stack(c_val)
+            if c_val.ndim > 1:
+                c_val = c_val.squeeze(-1)
+                
+        if c_val.numel() != self.m:
+            raise ValueError(f"Number of elements in c_val must be equal to m={self.m}, got {c_val.numel()}")
+        G = []
+
+        for group in self.param_groups:
+            params: list[Tensor] = []
+            grads: list[Tensor] = []
+            c_grads: list[Tensor] = []
+            smoothing: list[Tensor] = []
+            exp_avgs: list[Tensor] = []
+            exp_avg_sqs: list[Tensor] = []
+            max_exp_avg_sqs: list[Tensor] = []
+            state_steps: list[Tensor] = []
+            lr = group["lr"]
+            amsgrad = group["amsgrad"]
+            
+            _ = self._init_group(
+                group,
+                params,
+                grads,
+                c_grads,
+                exp_avgs,
+                exp_avg_sqs,
+                max_exp_avg_sqs,
+                state_steps,
+                smoothing
+            )
+
+            for i, param in enumerate(params):
+                ### calculate Lagrange f-n gradient (G) ###
+
+                # stack list of grads w.r.t. constraints to get
+                # tensor of shape (*param.shape, m)
+                l_term_grad = 0
+                aug_term_grad = 0
+                # if c_grads[i] is not None:
+                for j, c_grad in enumerate(c_grads[i]):
+                    if c_grad is None:
+                        continue
+                    l_term_grad += c_grad * self._dual_vars[j]
+                    aug_term_grad += c_grad * c_val[j]
+
+                G_i = (
+                    grads[i]
+                    + l_term_grad
+                    + self.rho * aug_term_grad
+                    + self.mu * (param - smoothing[i])
+                )
+                
+                G.append(G_i)
+                
+                exp_avg = exp_avgs[i]
+                exp_avg_sq = exp_avg_sqs[i]
+                step_t = state_steps[i]
+                step_t += 1
+                beta1 = self.beta1
+                beta2 = self.beta2
+                eps = self.eps
+                
+                exp_avg.lerp_(G_i, 1 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(G_i, G_i, value=1 - beta2)
+                
+
+                smoothing[i].add_(param - smoothing[i], alpha=self.beta)
+
+                # param.add_(G_i, alpha=-lr)
+                
+                bias_correction1 = 1 - beta1**step_t
+                bias_correction2 = 1 - beta2**step_t
+
+                step_size = lr / bias_correction1
+
+                bias_correction2_sqrt = bias_correction2**0.5
+
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])
+
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = (max_exp_avg_sqs[i].sqrt() / bias_correction2_sqrt).add_(eps)
+                else:
+                    denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
+
+                param.addcdiv_(exp_avg, denom, value=-step_size)
+
+                ## PROJECT (keep in mind we do layer by layer)
+                ## add slack variables to params in constructor?
+
+                c_grads[i].clear()
+        
+        self.c_vals.clear()
+        return G

humancompatible_train-0.1.3/humancompatible/train/algorithms/test/test_ssl_alm_adam.py

@@ -0,0 +1,146 @@
+import unittest
+import torch
+from torch import Tensor
+from humancompatible.train.algorithms import SSLALM_Adam
+
+class TestSSLALMAdam(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_Adam(
+            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,
+            beta1=0.9,
+            beta2=0.999
+        )
+
+    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}
+        self.optimizer._dual_vars = torch.ones(2)
+        
+        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.state[p]["step"] = 0
+            self.optimizer.state[p]["exp_avg"] = torch.ones_like(p)
+            self.optimizer.state[p]["exp_avg_sq"] = torch.ones_like(p)
+        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):
+            ## assert correct Lagrange f-n gradient
+            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)])
+                    )
+                )
+            )
+            # assert correct update of params
+            beta1 = 0.9
+            beta2 = 0.999
+            fm = beta1*torch.ones_like(p) + (1-beta1)*G[i]
+            sm = beta2*torch.ones_like(p) + (1-beta2)*(torch.pow(G[i],2))
+            fm_bc = fm/(1-beta1)
+            sm_bc = sm/(1-beta2)
+
+            self.assertTrue(
+                torch.equal(
+                    p,
+                    p_pre_step[p] - 0.01*fm_bc/(sm_bc.sqrt() + 1e-8)
+                )
+            )
+            # assert correct update of exp_avg and exp_avg_sq
+            self.assertTrue(
+                torch.equal(
+                    self.optimizer.state[p]['exp_avg'],
+                    fm
+                )
+            )
+            self.assertTrue(
+                torch.equal(
+                    self.optimizer.state[p]['exp_avg_sq'],
+                    sm
+                )
+            )
+
+    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_Adam(self.params, m=self.m, lr=-0.01)
+        with self.assertRaises(ValueError):
+            SSLALM_Adam(self.params, m=self.m, dual_lr=-0.01)
+        with self.assertRaises(ValueError):
+            SSLALM_Adam(self.params, m=self.m, init_dual_vars=torch.tensor([1.0]))
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: humancompatible-train
-Version: 0.1.2
+Version: 0.1.3
 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,6 +10,13 @@ Description-Content-Type: text/markdown
 License-File: LICENCE.txt
 Requires-Dist: torch
 Requires-Dist: numpy
+Provides-Extra: examples
+Requires-Dist: ipykernel; extra == "examples"
+Requires-Dist: ipympl; extra == "examples"
+Requires-Dist: fairret; extra == "examples"
+Requires-Dist: folktables; extra == "examples"
+Requires-Dist: scikit-learn; extra == "examples"
+Requires-Dist: matplotlib; extra == "examples"
 Provides-Extra: benchmark
 Requires-Dist: fairret; extra == "benchmark"
 Requires-Dist: matplotlib; extra == "benchmark"
@@ -31,22 +38,23 @@ The toolkit implements algorithms for constrained training of neural networks ba
 <!-- , 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. [Extending the toolkit](#extending-the-toolkit) 
+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
 requests, please file a
-[Github issue](https://github.com/humancompatible/train/issues). 
+[Github issue](https://github.com/humancompatible/train/issues).
 
 ## Installation
 
 Use
 
-```
+```bash
 pip install humancompatible-train
 ```
 
@@ -56,11 +64,11 @@ The only dependencies of this package are `numpy` and `torch`.
 
 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.
+The algorithms follow the `dual_step()` - `step()` framework: taking inspiration from PyTorch, the `dual_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:
 
-```
+```python
 for inputs, labels in dataloader:
   # inference
   outputs = model(inputs)
@@ -79,6 +87,7 @@ for inputs, labels in dataloader:
 ```
 
 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.
 
@@ -86,6 +95,12 @@ Our idea is to
 
 You are invited to check out our new API presented in notebooks in the `examples` folder.
 
+The example notebooks have additional dependencies, such as `fairret`. To install those, run
+
+```
+pip install humancompatible-train[examples]
+```
+
 *The legacy API used for the benchmark is presented in `examples/_old_/algorithm_demo.ipynb` and `examples/_old_/constraint_demo.ipynb`.*
 
 ## Extending the toolkit
@@ -99,21 +114,27 @@ You are invited to check out our new API presented in notebooks in the `examples
 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
 
 **bash** (Linux)
+
 ```
 python3.11 -m venv fairbenchenv
 source fairbenchenv/bin/activate
 ```
+
 **cmd** (Windows)
+
 ```
 python -m venv fairbenchenv
 fairbenchenv\Scripts\activate.bat
 ```
+
 2. Install from source.
+
 ```
 git clone https://github.com/humancompatible/train.git
 cd train
@@ -122,11 +143,12 @@ pip install .
 ```
 
 If you wish to edit the code of the algorithms, install as an editable package:
+
 ```
 pip install -e .
 ```
 
-__Warning__: it is recommended to use Stochastic Ghost with the mkl-accelerated version of the scipy package with Stochastic Ghost; to install it, run
+**Warning**: it is recommended to use Stochastic Ghost with the mkl-accelerated version of the scipy package with Stochastic Ghost; to install it, run
 
 ```pip install --force-reinstall -i https://software.repos.intel.com/python/pypi scipy```
 
@@ -140,12 +162,14 @@ after installing requirements.txt; otherwise, the algorithm will run slower. How
 ### Running the algorithms
 
 The benchmark comprises the following algorithms:
+
 - Stochastic Ghost [[2]](#2),
 - SSL-ALM [[3]](#3),
 - Stochastic Switching Subgradient [[4]](#4).
 
 To reproduce the experiments of the paper, run the following:
-``` python
+
+```
 cd experiments
 python run_folktables.py data=folktables alg=sslalm
 python run_folktables.py data=folktables alg=alm
@@ -154,23 +178,25 @@ python run_folktables.py data=folktables alg=ssg
 python run_folktables.py data=folktables alg=sgd     # baseline, no fairness
 python run_folktables.py data=folktables alg=fairret # baseline, fairness with regularizer
 ```
+
 Each command will start 10 runs of the `alg`, 30 seconds each.
 The results will be saved to `experiments/utils/saved_models` and `experiments/utils/exp_results`.
 <!-- In the repository, we include the configuration needed to reproduce the experiments in the paper. To do so, go to `experiments` and run `python run_folktables.py data=folktables alg=sslalm`. -->
 <!-- Repeat for the other algorithms by changing the `alg` parameter. -->
 
-This repository uses [Hydra](https://hydra.cc/) to manage parameters; see `experiments/conf` for configuration files. 
-* To change the parameters of the experiment, such as the number of runs for each algorithm, run time, the dataset used (*note: for now supports only Folktables*) - use `experiment.yaml`. 
-* To change the dataset settings - such as file location - or do dataset-specific adjustments - such as the configuration of the protected attributes - use `data/{dataset_name}.yaml`
-* To change algorithm hyperparameters, use `alg/{algorithm_name}.yaml`.
-* To change constraint hyperparameters, use `constraint/{constraint_name}.yaml`
+This repository uses [Hydra](https://hydra.cc/) to manage parameters; see `experiments/conf` for configuration files.
+
+- To change the parameters of the experiment, such as the number of runs for each algorithm, run time, the dataset used (*note: for now supports only Folktables*) - use `experiment.yaml`.
+- To change the dataset settings - such as file location - or do dataset-specific adjustments - such as the configuration of the protected attributes - use `data/{dataset_name}.yaml`
+- To change algorithm hyperparameters, use `alg/{algorithm_name}.yaml`.
+- To change constraint hyperparameters, use `constraint/{constraint_name}.yaml`
 
 <!-- ; it is installed as one of the dependencies. -->
 <!-- To learn more about using Hydra, please check out the [official tutorial](https://hydra.cc/docs/tutorials/basic/your_first_app). -->
 
 ### 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.
 
+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.
 
 ## License and terms of use
 
@@ -180,7 +206,7 @@ The benchmark part of the package relies on the Folktables package, provided und
 It provides code to download data from the American Community Survey
 (ACS) Public Use Microdata Sample (PUMS) files managed by the US Census Bureau.
 The data itself is governed by the terms of use provided by the Census Bureau.
-For more information, see https://www.census.gov/data/developers/about/terms-of-service.html
+For more information, see <https://www.census.gov/data/developers/about/terms-of-service.html>
 
 <!-- ## Cite this work -->
 
@@ -206,7 +232,7 @@ For more information, see https://www.census.gov/data/developers/about/terms-of-
 
 If you use this work, we encourage you to cite [our paper](https://arxiv.org/abs/2507.04033),
 
-```
+```bibtex
 @misc{kliachkin2025benchmarkingstochasticapproximationalgorithms,
       title={Benchmarking Stochastic Approximation Algorithms for Fairness-Constrained Training of Deep Neural Networks}, 
       author={Andrii Kliachkin and Jana Lepšová and Gilles Bareilles and Jakub Mareček},
@@ -218,14 +244,14 @@ If you use this work, we encourage you to cite [our paper](https://arxiv.org/abs
 }
 ```
 
-<a id="1">[1]</a> 
+<a id="1">[1]</a>
 Ding, Hardt & Miller et al. (2021) Retiring Adult: New Datasets for Fair Machine Learning, Curran Associates, Inc..
 
-<a id="2">[2]</a> 
+<a id="2">[2]</a>
 Facchinei & Kungurtsev (2023) Stochastic Approximation for Expectation Objective and Expectation Inequality-Constrained Nonconvex Optimization, arXiv.
 
-<a id="3">[3]</a> 
+<a id="3">[3]</a>
 Huang, Zhang & Alacaoglu (2025) Stochastic Smoothed Primal-Dual Algorithms for Nonconvex Optimization with Linear Inequality Constraints, arXiv.
 
-<a id="4">[4]</a> 
+<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.2 → humancompatible_train-0.1.3}/humancompatible_train.egg-info/SOURCES.txt

@@ -10,9 +10,11 @@ humancompatible/__init__.py
 humancompatible/train/__init__.py
 humancompatible/train/algorithms/__init__.py
 humancompatible/train/algorithms/ssl_alm.py
+humancompatible/train/algorithms/ssl_alm_adam.py
 humancompatible/train/algorithms/ssw.py
 humancompatible/train/algorithms/test/__init__.py
 humancompatible/train/algorithms/test/test_ssl_alm.py
+humancompatible/train/algorithms/test/test_ssl_alm_adam.py
 humancompatible/train/algorithms/test/test_ssw.py
 humancompatible/train/fairness/__init__.py
 humancompatible/train/fairness/constraints/__init__.py

{humancompatible_train-0.1.2 → humancompatible_train-0.1.3}/humancompatible_train.egg-info/requires.txt

@@ -10,6 +10,14 @@ pot
 hydra
 omegaconf
 
+[examples]
+ipykernel
+ipympl
+fairret
+folktables
+scikit-learn
+matplotlib
+
 [ghost]
 qpsolvers
 scipy

{humancompatible_train-0.1.2 → humancompatible_train-0.1.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "humancompatible-train"
-version = "0.1.2"
+version = "0.1.3"
 dependencies = [
     "torch",
     "numpy",
@@ -24,5 +24,6 @@ readme = "README.md"
 
 
 [project.optional-dependencies]
+examples = ["ipykernel", "ipympl", "fairret", "folktables", "scikit-learn", "matplotlib"]
 benchmark = ["fairret", "matplotlib", "pandas", "folktables", "pot", "hydra", "omegaconf"]
 ghost = ["qpsolvers", "scipy"]