humancompatible-train 0.1.0__py3-none-any.whl

humancompatible/train/algorithms/torch/ssl_alm.py

@@ -0,0 +1,212 @@
+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):
+    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,
+        *,
+        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,
+            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.mu = mu
+        self.c_vals: list[Union[float, Tensor]] = []
+        self._c_val_average = [None]
+        # 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, 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"] = []
+
+            smoothing.append(state.get("smoothing"))
+
+            grads.append(p.grad)
+            c_grads.append(state.get("c_grad"))
+
+        return has_sparse_grad
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+
+    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] = []
+            _ = self._init_group(group, params, grads, c_grads, 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] = []
+            lr = group["lr"]
+            _ = self._init_group(group, params, grads, c_grads, 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)
+
+                smoothing[i].add_(param - smoothing[i], alpha=self.beta)
+
+                param.add_(G_i, alpha=-lr)
+
+                ## 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/algorithms/torch/ssw.py

@@ -0,0 +1,155 @@
+from typing import Iterable, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch.optim.optimizer import Optimizer, _use_grad_for_differentiable
+
+# def project_fn(x, m):
+#     for i in range(1, m + 1):
+#         if x[-i] < 0:
+#             x[-i] = 0
+#     return x
+
+
+def _dual_step_func(dual_var, lr, cval):
+    return dual_var + lr * cval
+
+
+# def step_fn(params, grads,)
+
+
+class SSG(Optimizer):
+    def __init__(
+        self,
+        params,
+        m: int = 1,
+        # constraint tolerance
+        # ctols: Union[
+        #     float, Tensor
+        # ],
+        # ctols_rule = "const",
+        # learning rate
+        lr: Union[float, Tensor] = 5e-2,
+        # learning rate decrease rule
+        # lr_rule = "const",
+        # constraint learning rate
+        dual_lr: Union[
+            float, Tensor
+        ] = 5e-2,  # keep as tensor for different learning rates for different constraints in the future? idk
+        *,
+        differentiable: bool = False,
+    ):
+        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 not (m == 1):
+            raise ValueError(f"Switching Subgradient does not support multiple constraints."
+                             "Consider taking the largest violation at each iteration.")
+        if differentiable:
+            raise NotImplementedError("SSw does not support differentiable")
+
+        defaults = dict(
+            lr=lr,
+            dual_lr=dual_lr,
+            differentiable=differentiable,
+        )
+
+        super().__init__(params, defaults)
+
+        # self.param_groups.append()
+
+        self.m = m
+        self.lr = lr
+        self.dual_lr = dual_lr
+        # self.lr_rule = lr_rule
+        # self.dual_lr_rule = dual_lr_rule
+        self.c_vals: list[Union[float, Tensor]] = []
+        # self.ctols = ctols
+        # 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
+
+    def _init_group(self, group, params, grads, c_grads):
+        # 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)
+
+            # Lazy state initialization
+            if len(state) == 0:
+                state["c_grad"] = []
+
+            grads.append(p.grad)
+            c_grads.append(state.get("c_grad"))
+
+        return has_sparse_grad
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+
+    def dual_step(self, i: int, c_val: Tensor = None):
+        r"""Save constraint gradient for weight update. To be called BEFORE :func:`step` in an iteration!
+
+        Args:
+            i (int): index of the constraint **(unused)**
+            c_val (Tensor): an estimate of the value of the constraint at which the gradient was computed **(unused)**
+        """
+        
+        if i > self.m:
+            raise ValueError("SSw does not support multiple constraints.")
+
+        # save constraint grad
+        for group in self.param_groups:
+            params: list[Tensor] = []
+            grads: list[Tensor] = []
+            c_grads: list[Tensor] = []
+            _ = self._init_group(group, params, grads, c_grads)
+
+            for p in group["params"]:
+                if p.grad is not None:
+                    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]):
+        r"""Perform an update of the primal parameters (network weights). 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`
+        """
+
+        # here assume c_val is a scalar
+
+        update_con = c_val > 0
+        
+        for group in self.param_groups:
+            params: list[Tensor] = []
+            grads: list[Tensor] = []
+            c_grads: list[Tensor] = []
+            lr = group["lr"]
+            _ = self._init_group(group, params, grads, c_grads)
+
+            for i, param in enumerate(params):
+
+                if update_con:
+                    param.add_(c_grads[i][0], alpha=-self.dual_lr)
+                else:
+                    param.add_(param.grad, alpha=-lr)
+
+                if c_grads[i] is not None:
+                    c_grads[i].clear()

humancompatible/train/algorithms/utils.py

@@ -0,0 +1,61 @@
+import torch
+
+
+def net_params_to_tensor(
+    net: torch.nn.Module, flatten=False, copy=False
+) -> torch.Tensor:
+    # flat_params = [ar.to_numpy(param) for param in net.parameters()]
+    if copy:
+        params = [param.detach().clone() for param in net.parameters()]
+    else:
+        params = [param for param in net.parameters()]
+
+    if flatten:
+        flat_params = [torch.flatten(param) for param in params]
+        return torch.concat(flat_params)
+
+    return params
+
+
+def check_same_sample(sample1, sample2):
+    s1w, s1nw = sample1[0], sample1[1]
+    s2w, s2nw = sample2[0], sample2[1]
+
+    s1wx, s1wy = s1w
+    s2wx, s2wy = s2w
+
+    s1nwx, s1nwy = s1nw
+    s2nwx, s2nwy = s2nw
+
+    return (
+        torch.all(s1wx == s2wx)
+        and torch.all(s1wy == s2wy)
+        and torch.all(s1nwx == s2nwx)
+        and torch.all(s1nwy == s2nwy)
+    )
+
+
+def net_grads_to_tensor(net, clip=False, flatten=True, device=None) -> torch.Tensor:
+    param_grads = []
+    if clip:
+        torch.nn.utils.clip_grad_norm_(net.parameters(), 0.5)
+    for param in net.parameters():
+        if param.grad is not None:
+            # Clone to avoid modifying the original tensor
+            device = param.grad.data.device if device is None else device
+            if flatten:
+                param_grads.append(param.grad.data.view(-1))
+            else:
+                param_grads.append(param.grad.data.to(device))
+    if flatten:
+        param_grads = torch.cat(param_grads)
+    return param_grads
+
+
+def _set_weights(net: torch.nn.Module, x):
+    start = 0
+    w = net_params_to_tensor(net, flatten=False, copy=False)
+    for i in range(len(w)):
+        end = start + w[i].numel()
+        w[i].set_(x[start:end].reshape(w[i].shape))
+        start = end

humancompatible/train/constraints/__init__.py

@@ -0,0 +1,11 @@
+from .constraint import FairnessConstraint
+from .constraint_fns import (
+    fairret_stat_equality,
+    ppv_equality,
+    acc_equality,
+    tpr_equality,
+    abs_loss_equality,
+    loss_equality,
+)
+
+__all__ = ["FairnessConstraint, loss_equality"]

humancompatible/train/constraints/constraint.py

@@ -0,0 +1,87 @@
+from typing import Callable, Iterable
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, SubsetRandomSampler
+
+
+def _make_dataloaders(dataset, group_indices, batch_size, device, drop_last, gen=None):
+    dataloaders = []
+    for idx in group_indices:
+        sampler = SubsetRandomSampler(idx, gen)
+        dataloaders.append(iter(DataLoader(dataset, batch_size, sampler=sampler, drop_last=True)))
+    return dataloaders
+
+
+class FairnessConstraint:
+    def __init__(
+        self,
+        dataset: torch.utils.data.Dataset,
+        group_indices: Iterable[Iterable[int]],
+        fn: Callable,
+        batch_size: int = None,
+        use_dataloaders=True,
+        device="cpu",
+        seed=None,
+        loader_drop_last=False,
+    ):
+        self.dataset = dataset
+        self.group_sets = [
+            torch.utils.data.Subset(dataset, idx) for idx in group_indices
+        ]
+        self._group_indices = group_indices
+        self.fn = fn
+        self._seed = seed
+        self._rng = np.random.default_rng(seed)
+        self._torch_rng = torch.manual_seed(seed) if seed is not None else torch.Generator(device=device)
+        self._device = device
+        self._drop_last = loader_drop_last
+        if batch_size is not None:
+            self._batch_size = batch_size
+            if use_dataloaders:
+                self.group_dataloaders = _make_dataloaders(
+                    dataset, group_indices, batch_size, device, gen=self._torch_rng, drop_last=loader_drop_last
+                )
+
+    def group_sizes(self):
+        return [len(group) for group in self.group_sets]
+
+    def eval(self, net, sample, **kwargs):
+        return self.fn(net, sample, **kwargs)
+
+    def sample_loader(self):
+        samples = []
+        for i, l in enumerate(self.group_dataloaders):
+            try:
+                sample = next(l)
+            except StopIteration:
+                sampler = SubsetRandomSampler(self._group_indices[i], self._torch_rng)
+                l = iter(DataLoader(self.dataset, self._batch_size, sampler=sampler, drop_last=self._drop_last))
+                sample = next(l)
+                self.group_dataloaders[i] = l
+                
+            samples.append(sample)
+        return samples
+
+    def sample_dataset(
+        self, N, rng: np.random.Generator = None, indices=None, return_indices=False
+    ):
+        if rng is None:
+            rng = self._rng
+
+        if indices is None:
+            indices = []
+            # returns len(group) points if N > len(group)
+            for group in self.group_sets:
+                indices.append(
+                    rng.choice(group.indices, N)
+                    if N < len(group)
+                    else rng.choice(group.indices, len(group))
+                )
+
+        sample = [self.dataset[indices[i]] for i, _ in enumerate(self.group_sets)]
+
+        if return_indices:
+            return sample, indices
+        else:
+            return sample

humancompatible/train/constraints/constraint_fns.py

@@ -0,0 +1,118 @@
+import torch
+from fairret.statistic import (
+    TruePositiveRate,
+    FalseNegativeFalsePositiveFraction,
+    Accuracy,
+)
+from fairret.loss import NormLoss
+
+
+def tpr_equality(_, net, c_data):
+    statistic = TruePositiveRate()
+    loss = NormLoss(statistic, p=2)
+
+    return fairret_stat_equality(net, c_data, loss)
+
+
+def ppv_equality(_, net, c_data):
+    statistic = FalseNegativeFalsePositiveFraction()
+    loss = NormLoss(statistic, p=2)
+
+    return fairret_stat_equality(net, c_data, loss)
+
+
+def acc_equality(_, net, c_data):
+    statistic = Accuracy()
+    loss = NormLoss(statistic, p=2)
+
+    return fairret_stat_equality(net, c_data, loss)
+
+
+def fairret_stat_equality(net, c_data, loss):
+    g1_inputs, g1_labels = c_data[0]
+    g2_inputs, g2_labels = c_data[1]
+
+    g1_outs = net(g1_inputs).squeeze()
+    g2_outs = net(g2_inputs).squeeze()
+    # if not (1 in a_labels or 1 in g2_labels):
+    #     return torch.tensor(0)
+
+    group_codes = [0] * len(g1_labels) + [1] * len(g2_labels)
+    group_codes = torch.tensor(
+        [[0.0, 1.0] if x == 1 else [1.0, 0.0] for x in group_codes]
+    )
+
+    return loss(
+        torch.concat([g1_outs, g2_outs]).unsqueeze(1),
+        group_codes,
+        torch.concat([g1_labels, g2_labels]).unsqueeze(1),
+    )
+
+
+def dummy(_, net, c_data):
+    r = torch.zeros(1)
+    r.grad = 0
+    return r
+
+
+def loss_equality(loss, net, c_data):
+    g1_inputs, g1_labels = c_data[0]
+    g2_inputs, g2_labels = c_data[1]
+    g1_outs = net(g1_inputs)
+    if g1_labels.ndim == 0:
+        g1_labels = g1_labels.reshape(1)
+        g2_labels = g2_labels.reshape(1)
+    if g1_labels.ndim < g1_outs.ndim:
+        g1_labels = g1_labels.unsqueeze(1)
+        g2_labels = g2_labels.unsqueeze(1)
+    g1_loss = loss(g1_outs, g1_labels)
+    g2_outs = net(g2_inputs)
+    g2_loss = loss(g2_outs, g2_labels)
+
+    val = g1_loss - g2_loss
+    return val
+
+
+def abs_loss_equality(loss, net, c_data):
+    g1_inputs, g1_labels = c_data[0]
+    g2_inputs, g2_labels = c_data[1]
+    g1_outs = net(g1_inputs)
+    if g1_labels.ndim == 0:
+        g1_labels = g1_labels.reshape(1)
+        g2_labels = g2_labels.reshape(1)
+    if g1_labels.ndim < g1_outs.ndim:
+        g1_labels = g1_labels.unsqueeze(1)
+        g2_labels = g2_labels.unsqueeze(1)
+    g1_loss = loss(g1_outs, g1_labels)
+    g2_outs = net(g2_inputs)
+    g2_loss = loss(g2_outs, g2_labels)
+
+    val = g1_loss - g2_loss
+    return torch.abs(val)
+
+
+def fairret_constr(loss, net, c_data):
+    g1_inputs, g1_labels = c_data[0]
+    g2_inputs, g2_labels = c_data[1]
+    g1_logits = net(g1_inputs)
+    g2_logits = net(g2_inputs)
+    g1_onehot = torch.tensor([[0.0, 1.0]] * len(g1_inputs))
+    g2_onehot = torch.tensor([[1.0, 0.0]] * len(g2_inputs))
+    logits = torch.concat([g1_logits, g2_logits])
+    sens = torch.vstack([g1_onehot, g2_onehot])
+    labels = torch.hstack([g1_labels, g2_labels]).unsqueeze(1)
+
+    return loss(logits, sens, label=labels)
+
+
+def fairret_pr_constr(loss, net, c_data):
+    g1_inputs, _ = c_data[0]
+    g2_inputs, _ = c_data[1]
+    g1_logits = net(g1_inputs)
+    g2_logits = net(g2_inputs)
+    g1_onehot = torch.tensor([[0.0, 1.0]] * len(g1_inputs))
+    g2_onehot = torch.tensor([[1.0, 0.0]] * len(g2_inputs))
+    logits = torch.concat([g1_logits, g2_logits])
+    sens = torch.vstack([g1_onehot, g2_onehot])
+
+    return loss(logits, sens)

humancompatible/train/fairness/constraints/__init__.py

@@ -0,0 +1,15 @@
+from .constraint import FairnessConstraint
+from .constraint_fns import (
+    fairret_stat_equality,
+    ppv_equality,
+    acc_equality,
+    tpr_equality,
+    abs_loss_equality,
+    loss_equality,
+    abs_diff_tpr,
+    abs_diff_fpr,
+    abs_max_dev_from_overall_tpr,
+    abs_diff_pr
+)
+
+__all__ = ["FairnessConstraint, loss_equality"]