From cb4fe94b96114cb8baa506cec19ccd1cd90575ac Mon Sep 17 00:00:00 2001 From: Alice Driessen Date: Mon, 13 Jan 2025 16:30:53 +0100 Subject: [PATCH 1/6] First commit --- .../neural/methods/conditional_monge_gap.py | 478 ++++++++++++++++++ .../conditional_perturbation_network.py | 117 +++++ 2 files changed, 595 insertions(+) create mode 100644 src/ott/neural/methods/conditional_monge_gap.py create mode 100644 src/ott/neural/networks/conditional_perturbation_network.py diff --git a/src/ott/neural/methods/conditional_monge_gap.py b/src/ott/neural/methods/conditional_monge_gap.py new file mode 100644 index 000000000..688876491 --- /dev/null +++ b/src/ott/neural/methods/conditional_monge_gap.py @@ -0,0 +1,478 @@ +import collections +import functools + +from pathlib import Path +from typing import ( + Any, + Callable, + Dict, + Iterator, + Literal, + Optional, + Sequence, + Tuple, + Type, + TypeVar, + Union, +) + +import flax.linen as nn +import jax +import jax.numpy as jnp +import optax +from flax.core import frozen_dict +from flax.training import train_state +from flax.training.orbax_utils import save_args_from_target +from jax.tree_util import tree_map +from orbax.checkpoint import PyTreeCheckpointer +from ott.geometry import costs, pointcloud +from ott.solvers import linear +from ott.solvers.linear import sinkhorn +from ott.neural.networks.conditional_perturbation_network import ( + ConditionalPerturbationNetwork, +) + +T = TypeVar("T", bound="ConditionalMongeGapEstimator") + + +def cmonge_gap_from_samples( + source: jnp.ndarray, + target: jnp.ndarray, + target_condition: jnp.ndarray, + source_condition: Optional[jnp.ndarray], + equal_conditions: bool = False, + cost_fn: Optional[costs.CostFn] = None, + epsilon: Optional[float] = None, + relative_epsilon: Optional[Literal["mean", "std"]] = None, + scale_cost: Union[float, Literal["mean", "max_cost", "median"]] = 1.0, + return_output: bool = False, + rng: Optional[jax.Array] = None, + **kwargs: Any, +) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: + r"""Monge gap, instantiated in terms of samples before / after applying map. + + .. math:: + \sum_{i=1}{K} \frac{1}{n} \sum_{i=1}^n c(x_i, y_i)) - + W_{c, \varepsilon}(\frac{1}{n}\sum_i \delta_{x_i}, + \frac{1}{n}\sum_i \delta_{y_i}) + + where :math:`W_{c, \varepsilon}` is an + :term:`entropy-regularized optimal transport` + cost, the :attr:`~ott.solvers.linear.sinkhorn.SinkhornOutput.ent_reg_cost`. + + Args: + source: samples from first measure, array of shape ``[n, d]``. + target: samples from second measure, array of shape ``[n, d]``. + target_condition: array indicating condition to which each target sample + belongs, `integer array of shape ``[n]``. + source_condition: array indicating condition to which each source sample + belongs, `integer array of shape ``[n]``. + If `equal_condition` is `None` and `source_condition` is `False` + per condition same number of source cells are sampled as there are + target cells. + equal_conditions: whether source and target samples come from the same + (order) of conditions. In this case `target_conditions` is used + for both. + cost_fn: a cost function between two points in dimension :math:`d`. + If :obj:`None`, :class:`~ott.geometry.costs.SqEuclidean` is used. + epsilon: Regularization parameter. See + :class:`~ott.geometry.pointcloud.PointCloud` + relative_epsilon: when `False`, the parameter ``epsilon`` specifies the + value of the entropic regularization parameter. When `True`, ``epsilon`` + refers to a fraction of the + :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is + computed adaptively using ``source`` and ``target`` points. + scale_cost: option to rescale the cost matrix. Implemented scalings are + 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be + given to rescale the cost such that ``cost_matrix /= scale_cost``. + return_output: boolean to also return the + :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. + rng: random key used for sampling source samples. + kwargs: holds the kwargs to instantiate the or + :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to + compute the regularized OT cost. + + Returns: + The average Monge gap value over all conditions and optionally the + list of Monge gap per condition and :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` + """ + key = jax.random.PRNGKey(rng) + cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn + all_losses = [] + all_outs = [] + for c in jnp.unique(target_condition): + key, _ = jax.random.split(key, 2) + c_target = target[target_condition == c] + if equal_conditions: + c_source = source[target_condition == c] + elif source_condition: + c_source = source[source_condition == c] + else: + c_source = jax.random.choice( + key, + len(source), + size=len(c_target), + replace=len(c_target) > len(source), + ) + geom = pointcloud.PointCloud( + x=c_source, + y=c_target, + cost_fn=cost_fn, + epsilon=epsilon, + relative_epsilon=relative_epsilon, + scale_cost=scale_cost, + ) + gt_displacement_cost = jnp.mean(jax.vmap(cost_fn)(source, target)) + out = linear.solve(geom=geom, **kwargs) + loss = gt_displacement_cost - out.ent_reg_cost + all_losses.append(loss) + if return_output: + all_outs.append(out) + loss = sum(all_losses) / len(all_losses) # average + + return (loss, out, all_losses) if return_output else loss + + +class ConditionalMongeGapEstimator: + r"""Monge Gap Estimator which optimizes over multiple conditions + + .. math:: + \text{min}_{\theta}\; \sum_{i=1}^{K} \Delta(T_\theta \sharp \mu, \theta) + + \lambda R(T_\theta \sharp \rho, \rho) + + Args: + dim_data: input dimensionality of data required for network init. + dim_cond: input dimensionality of condition embedding, + required for network init + model: network architecture for map :math:`T`, + should be a `ConditionalPerturbationNetwork`. + optimizer: optimizer function for map :math:`T`. + fitting_loss: function that outputs a fitting loss :math:`\Delta` between + two families of points, as well as any log object. + regularizer: function that outputs a score from two families of points, + here assumed to be of the same size, as well as any log object. + regularizer_strength: strength of the :attr:`regularizer`. + num_train_iters: number of total training iterations. + logging: option to return logs. + valid_freq: frequency with training and validation are logged. + rng: random key used for seeding for network initializations. + """ + + def __init__( + self, + dim_data: int, + dim_cond: int, + model: ConditionalPerturbationNetwork, + optimizer: Optional[optax.OptState] = None, + fitting_loss: Optional[ + Callable[[jnp.ndarray, jnp.ndarray], Tuple[float, Optional[Any]]] + ] = None, + regularizer: Optional[ + Callable[[jnp.ndarray, jnp.ndarray], Tuple[float, Optional[Any]]] + ] = None, + regularizer_strength: Union[float, Sequence[float]] = 1.0, + num_train_iters: int = 10_000, + logging: bool = False, + valid_freq: int = 500, + grad_acc_steps: int = 1, + rng: Optional[jax.Array] = None, + ) -> None: + self._fitting_loss = fitting_loss + self._regularizer = regularizer + self.num_train_iters = self.num_train_iters + self.grad_acc_steps = grad_acc_steps + # Can use either a fixed strength, or generalize to a schedule. + self.regularizer_strength = jnp.repeat( + jnp.atleast_2d(regularizer_strength), + num_train_iters, + total_repeat_length=num_train_iters, + axis=0, + ).ravel() + self.logging = logging + self.valid_freq = valid_freq + self.rng = jax.random.PRNGKey(rng) + + # set default optimizer + if optimizer is None: + optimizer = optax.adam(learning_rate=0.001) + + # setup training + self.setup(dim_data, dim_cond, model, optimizer) + + def setup( + self, + dim_data: int, + dim_cond: int, + neural_net: ConditionalPerturbationNetwork, + optimizer: optax.OptState, + ): + """Setup all components required to train the network""" + + # neural network + self.rng, rng = jax.random.split(self.key, 2) + self.state_neural_net = neural_net.create_train_state( + self.rng, optimizer, dim_data, dim_cond + ) + + # step function + self.step_fn = self._get_step_fn() + + @property + def regularizer(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: + """Regularizer added to the fitting loss. + + Can be, e.g. the + :func:`~ott.neural.methods.monge_gap.monge_gap_from_samples`. + If no regularizer is passed for solver instantiation, + or regularization weight :attr:`regularizer_strength` is 0, + return 0 by default along with an empty set of log values. + """ # noqa: E501 + if self._regularizer is not None: + return self._regularizer + return lambda *_, **__: (0.0, None) + + @property + def fitting_loss(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: + """Fitting loss to fit the marginal constraint. + + Can be, e.g. :func:`~ott.tools.sinkhorn_divergence.sinkdiv`. + If no fitting_loss is passed for solver instantiation, return 0 by default, + and no log values. + """ + if self._fitting_loss is not None: + return self._fitting_loss + return lambda *_, **__: (0.0, None) + + def _generate_batch( + self, loader_source, loader_target + ) -> Dict[str, jnp.ndarray]: + """Generate a batch of condition and samples.""" + return { + "source": next(loader_source), + "target": next(loader_target), + } + + def train_map_estimator( + self, + trainloader_source: Iterator[jnp.ndarray], + trainloader_target: Iterator[jnp.ndarray], + validloader_source: Iterator[jnp.ndarray], + validloader_target: Iterator[jnp.ndarray], + ): + # define logs + logs = collections.defaultdict(lambda: collections.defaultdict(list)) + + # try to display training progress with tqdm + try: + from tqdm import trange + + tbar = trange(self.num_train_iters, leave=True) + except ImportError: + tbar = range(self.num_train_iters) + + grads = tree_map(jnp.zeros_like, self.state_neural_net.params) + for step in tbar: + # update step + is_logging_step = self.logging and ( + (step % self.valid_freq == 0) + or (step == self.num_train_iters - 1) + ) + is_gradient_acc_step = (step + 1) % self.grad_acc_steps == 0 + train_batch, condition = self._generate_batch( + trainloader_source, trainloader_target + ) + valid_batch, _ = ( + None + if not is_logging_step + else self._generate_batch( + validloader_source, validloader_target + ) + ) + + self.state_neural_net, grads, current_logs = self.step_fn( + self.state_neural_net, + grads=grads, + train_batch=train_batch, + valid_batch=valid_batch, + is_logging_step=is_logging_step, + is_gradient_acc_step=is_gradient_acc_step, + ) + + # store and print metrics if logging step + if is_logging_step: + for log_key in current_logs: + for metric_key in current_logs[log_key]: + logs[log_key][metric_key].append( + current_logs[log_key][metric_key] + ) + + # update the tqdm bar if tqdm is available + if not isinstance(tbar, range): + reg_msg = ( + "NA" + if current_logs["eval"]["regularizer"] == 0.0 + else f"{current_logs['eval']['regularizer']:.4f}" + ) + postfix_str = ( + f"fitting_loss: {current_logs['eval']['fitting_loss']:.4f}, " + f"regularizer: {reg_msg} ," + f"total: {current_logs['eval']['total_loss']:.4f}" + ) + tbar.set_postfix_str(postfix_str) + + return self.state_neural_net, logs + + def _get_step_fn(self) -> Callable: + """Create a one step training and evaluation function.""" + + def loss_fn( + params: frozen_dict.FrozenDict, + apply_fn: Callable, + batch: Dict[str, jnp.ndarray], + ) -> Tuple[float, Dict[str, float]]: + """Loss function.""" + # map samples with the fitted map + mapped_samples = apply_fn( + {"params": params}, + batch["source"]["X"], + batch["target"]["c_embed"], + ) + + # compute the loss + val_fitting_loss, log_fitting_loss = self.fitting_loss( + batch["target"], mapped_samples + ) + val_regularizer, log_regularizer = self.regularizer( + batch["source"], mapped_samples + ) + val_tot_loss, log_regularizer = val_fitting_loss + val_regularizer + + # store training logs + loss_logs = { + "total_loss": val_tot_loss, + "fitting_loss": val_fitting_loss, + "regularizer": val_regularizer, + "log_regularizer": log_regularizer, + "log_fitting": log_fitting_loss, + } + + return val_tot_loss, loss_logs + + @functools.partial(jax.jit, static_argnums=[4, 5]) + def step_fn( + state_neural_net: train_state.TrainState, + grads: frozen_dict.FrozenDict, + train_batch: Dict[str, jnp.ndarray], + valid_batch: Optional[Dict[str, jnp.ndarray]] = None, + is_logging_step: bool = False, + is_gradient_acc_step: bool = False, + ) -> Tuple[ + train_state.TrainState, frozen_dict.FrozenDict, Dict[str, float] + ]: + """Step function.""" + # compute loss and gradients + grad_fn = jax.value_and_grad(loss_fn, argnums=0, has_aux=True) + (_, current_train_logs), step_grads = grad_fn( + state_neural_net.params, + state_neural_net.apply_fn, + train_batch, + ) + # Accumulate gradients + grads = tree_map(lambda g, step_g: g + step_g, grads, step_grads) + + # logging step + current_logs = {"train": current_train_logs, "eval": {}} + if is_logging_step: + _, current_eval_logs = loss_fn( + params=state_neural_net.params, + apply_fn=state_neural_net.apply_fn, + batch=valid_batch, + ) + current_logs["eval"] = current_eval_logs + + # update state + if is_gradient_acc_step: + state_neural_net = state_neural_net.apply_gradients( + grads=tree_map(lambda g: g / self.grad_acc_steps, grads) + ) + # Reset gradients + grads = tree_map(jnp.zeros_like, grads) + + return state_neural_net, grads, current_logs + + return step_fn + + def transport(self, x, c): + return self.state_neural_net.apply_fn( + {"params": self.state_neural_net.params}, x, c + ) + + @property + def model(self) -> nn.Module: + return self.state_neural_net + + @model.setter + def model(self, value: nn.Module): + """Setter for the model to be checkpointed.""" + self.state_neural_net = value + + def save_checkpoint(self, path: Optional[Path] = None) -> None: + """Abstract method for saving model parameters to a pickle file. + + Args: + path: Path where the checkpoint should be saved. Defaults to None in which case + it is retrieved from config. + config: The model training configuration with a `checkpointing_path` field. + Defaults to None. + NOTE: If `config` and `path` are both not None, `path` takes preference. + """ + if path is None: + raise ValueError( + "Checkpoint cannot be saved. Provide a checkpoint save path" + ) + try: + checkpointer = PyTreeCheckpointer() + save_args = save_args_from_target(self.model) + checkpointer.save(path, self.model, save_args=save_args, force=True) + except Exception as e: + raise Exception(f"Error in saving checkpoint to {path}: {e}") + + @classmethod + def load_checkpoint( + cls: Type[T], + ckpt_path: Path = None, + *args, + **kwargs, + ) -> T: + """ + Loading a model from a checkpoint + + Args: + cls: Class object to be created. + ckpt_path: Optional path from where checkpoint is restored. + Defaults to None, in that case inferred from config. + *args: args normally given to `ConditionalMongeGapEstimator` + + Returns: + Class object with restored weights. + """ + try: + out_class = cls( + *args, + **kwargs, + ) + checkpointer = PyTreeCheckpointer() + out_class.model = checkpointer.restore( + ckpt_path, item=out_class.model + ) + return out_class + except Exception as e: + raise Exception( + f"Failed to load checkpoin from {ckpt_path}: {e}\nAre you sure" + "checkpoint was saved and correct path is provided?" + ) + + +# Add condition embedding to dataloader --> Users +# How to sample condition...? +# Optim/embedding/regularization factory? diff --git a/src/ott/neural/networks/conditional_perturbation_network.py b/src/ott/neural/networks/conditional_perturbation_network.py new file mode 100644 index 000000000..81fefff76 --- /dev/null +++ b/src/ott/neural/networks/conditional_perturbation_network.py @@ -0,0 +1,117 @@ +from typing import Any, Callable, Iterable, Sequence, Tuple + +import flax.linen as nn +import jax.numpy as jnp +import optax +from ott.neural.networks.potentials import ( + BasePotential, + PotentialTrainState, +) + + +class ConditionalPerturbationNetwork(BasePotential): + dim_hidden: Sequence[int] = None + dim_data: int = None + dim_cond: int = None # Full dimension of all context variables concatenated + # Same length as context_entity_bonds if embed_cond_equal is False + # (if True, first item is size of deep set layer, rest is ignored) + dim_cond_map: Iterable[int] = (50,) + act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.gelu + is_potential: bool = False + layer_norm: bool = False + embed_cond_equal: bool = ( + False # Whether all context variables should be treated as set or not + ) + context_entity_bonds: Iterable[Tuple[int, int]] = ( + (0, 10), + (0, 11), + ) # Start/stop index per modality + num_contexts: int = 2 + + @nn.compact + def __call__( + self, x: jnp.ndarray, c: jnp.ndarray + ) -> jnp.ndarray: # noqa: D102 + """ + Args: + x (jnp.ndarray): The input data of shape bs x dim_data + c (jnp.ndarray): The context of shape bs x dim_cond with + possibly different modalities + concatenated, as can be specified via context_entity_bonds. + + Returns: + jnp.ndarray: _description_ + """ + n_input = x.shape[-1] + + # Chunk the inputs + contexts = [ + c[:, e[0] : e[1]] + for i, e in enumerate(self.context_entity_bonds) + if i < self.num_contexts + ] + + if not self.embed_cond_equal: + # Each context is processed by a different layer, + # good for combining modalities + assert len(self.context_entity_bonds) == len(self.dim_cond_map), ( + "Length of context entity bonds and context map sizes have to " + f"match: {self.context_entity_bonds} != {self.dim_cond_map}" + ) + + layers = [ + nn.Dense(self.dim_cond_map[i], use_bias=True) + for i in range(len(contexts)) + ] + embeddings = [ + self.act_fn(layers[i](context)) + for i, context in enumerate(contexts) + ] + cond_embedding = jnp.concatenate(embeddings, axis=1) + else: + # We can use any number of contexts from the same modality, + # via a permutation-invariant deep set layer. + sizes = [c.shape[-1] for c in contexts] + if not len(set(sizes)) == 1: + raise ValueError( + "For embedding a set, all contexts need same length ," + f"not {sizes}" + ) + layer = nn.Dense(self.dim_cond_map[0], use_bias=True) + embeddings = [self.act_fn(layer(context)) for context in contexts] + # Average along stacked dimension + # (alternatives like summing are possible) + cond_embedding = jnp.mean(jnp.stack(embeddings), axis=0) + + z = jnp.concatenate((x, cond_embedding), axis=1) + if self.layer_norm: + n = nn.LayerNorm() + z = n(z) + + for n_hidden in self.dim_hidden: + wx = nn.Dense(n_hidden, use_bias=True) + z = self.act_fn(wx(z)) + wx = nn.Dense(n_input, use_bias=True) + + return x + wx(z) + + def create_train_state( + self, + rng: jnp.ndarray, + optimizer: optax.OptState, + dim_data: int, + dim_cond: int, + **kwargs: Any, + ) -> PotentialTrainState: + """Create initial `TrainState`.""" + c = jnp.ones((1, dim_cond)) # (n_batch, embed_dim) + x = jnp.ones((1, dim_data)) # (n_batch, data_dim) + params = self.init(rng, x=x, c=c)["params"] + return PotentialTrainState.create( + apply_fn=self.apply, + params=params, + tx=optimizer, + potential_value_fn=self.potential_value_fn, + potential_gradient_fn=self.potential_gradient_fn, + **kwargs, + ) From 294bc0d879a4b14ad8842b81e034c516bfef848b Mon Sep 17 00:00:00 2001 From: Alice Driessen Date: Mon, 13 Jan 2025 16:34:40 +0100 Subject: [PATCH 2/6] Added notes on dataloader structure for training --- src/ott/neural/methods/conditional_monge_gap.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/ott/neural/methods/conditional_monge_gap.py b/src/ott/neural/methods/conditional_monge_gap.py index 688876491..56951ff66 100644 --- a/src/ott/neural/methods/conditional_monge_gap.py +++ b/src/ott/neural/methods/conditional_monge_gap.py @@ -259,6 +259,9 @@ def train_map_estimator( validloader_source: Iterator[jnp.ndarray], validloader_target: Iterator[jnp.ndarray], ): + """The dataloaders should return a dict with key `X`. + The target dataloaders should additionally include a key + `c_embed`, which has the embedded condition in `dim_cond`.""" # define logs logs = collections.defaultdict(lambda: collections.defaultdict(list)) @@ -340,10 +343,10 @@ def loss_fn( # compute the loss val_fitting_loss, log_fitting_loss = self.fitting_loss( - batch["target"], mapped_samples + batch["target"]["X"], mapped_samples ) val_regularizer, log_regularizer = self.regularizer( - batch["source"], mapped_samples + batch["source"]["X"], mapped_samples ) val_tot_loss, log_regularizer = val_fitting_loss + val_regularizer From b94fef69af38344629bc60daea056e8a04dcedab Mon Sep 17 00:00:00 2001 From: Alice Driessen Date: Mon, 20 Jan 2025 14:35:52 +0100 Subject: [PATCH 3/6] Added conditional monge gap from samples --- .../neural/methods/conditional_monge_gap.py | 438 +----------------- 1 file changed, 24 insertions(+), 414 deletions(-) diff --git a/src/ott/neural/methods/conditional_monge_gap.py b/src/ott/neural/methods/conditional_monge_gap.py index 56951ff66..3b52a366e 100644 --- a/src/ott/neural/methods/conditional_monge_gap.py +++ b/src/ott/neural/methods/conditional_monge_gap.py @@ -7,7 +7,6 @@ Callable, Dict, Iterator, - Literal, Optional, Sequence, Tuple, @@ -25,8 +24,7 @@ from flax.training.orbax_utils import save_args_from_target from jax.tree_util import tree_map from orbax.checkpoint import PyTreeCheckpointer -from ott.geometry import costs, pointcloud -from ott.solvers import linear +from ott.neural.methods.monge_gap import monge_gap_from_samples from ott.solvers.linear import sinkhorn from ott.neural.networks.conditional_perturbation_network import ( ConditionalPerturbationNetwork, @@ -38,15 +36,8 @@ def cmonge_gap_from_samples( source: jnp.ndarray, target: jnp.ndarray, - target_condition: jnp.ndarray, - source_condition: Optional[jnp.ndarray], - equal_conditions: bool = False, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - relative_epsilon: Optional[Literal["mean", "std"]] = None, - scale_cost: Union[float, Literal["mean", "max_cost", "median"]] = 1.0, + condition: jnp.ndarray, return_output: bool = False, - rng: Optional[jax.Array] = None, **kwargs: Any, ) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: r"""Monge gap, instantiated in terms of samples before / after applying map. @@ -63,419 +54,38 @@ def cmonge_gap_from_samples( Args: source: samples from first measure, array of shape ``[n, d]``. target: samples from second measure, array of shape ``[n, d]``. - target_condition: array indicating condition to which each target sample - belongs, `integer array of shape ``[n]``. - source_condition: array indicating condition to which each source sample - belongs, `integer array of shape ``[n]``. - If `equal_condition` is `None` and `source_condition` is `False` - per condition same number of source cells are sampled as there are - target cells. - equal_conditions: whether source and target samples come from the same - (order) of conditions. In this case `target_conditions` is used - for both. - cost_fn: a cost function between two points in dimension :math:`d`. - If :obj:`None`, :class:`~ott.geometry.costs.SqEuclidean` is used. - epsilon: Regularization parameter. See - :class:`~ott.geometry.pointcloud.PointCloud` - relative_epsilon: when `False`, the parameter ``epsilon`` specifies the - value of the entropic regularization parameter. When `True`, ``epsilon`` - refers to a fraction of the - :attr:`~ott.geometry.pointcloud.PointCloud.mean_cost_matrix`, which is - computed adaptively using ``source`` and ``target`` points. - scale_cost: option to rescale the cost matrix. Implemented scalings are - 'median', 'mean' and 'max_cost'. Alternatively, a float factor can be - given to rescale the cost such that ``cost_matrix /= scale_cost``. + condition: array indicating condition for each source-target sample + `integer array of shape ``[n]``. return_output: boolean to also return the :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput`. - rng: random key used for sampling source samples. - kwargs: holds the kwargs to instantiate the or - :class:`~ott.solvers.linear.sinkhorn.Sinkhorn` solver to - compute the regularized OT cost. + kwargs: holds the kwargs to the function + :function:`~ott.neural.methods.monge_gap.monge_gap_from_samples` Returns: The average Monge gap value over all conditions and optionally the list of Monge gap per condition and :class:`~ott.solvers.linear.sinkhorn.SinkhornOutput` """ - key = jax.random.PRNGKey(rng) - cost_fn = costs.SqEuclidean() if cost_fn is None else cost_fn - all_losses = [] + all_gaps = [] all_outs = [] - for c in jnp.unique(target_condition): - key, _ = jax.random.split(key, 2) - c_target = target[target_condition == c] - if equal_conditions: - c_source = source[target_condition == c] - elif source_condition: - c_source = source[source_condition == c] - else: - c_source = jax.random.choice( - key, - len(source), - size=len(c_target), - replace=len(c_target) > len(source), - ) - geom = pointcloud.PointCloud( - x=c_source, - y=c_target, - cost_fn=cost_fn, - epsilon=epsilon, - relative_epsilon=relative_epsilon, - scale_cost=scale_cost, - ) - gt_displacement_cost = jnp.mean(jax.vmap(cost_fn)(source, target)) - out = linear.solve(geom=geom, **kwargs) - loss = gt_displacement_cost - out.ent_reg_cost - all_losses.append(loss) - if return_output: - all_outs.append(out) - loss = sum(all_losses) / len(all_losses) # average - - return (loss, out, all_losses) if return_output else loss - - -class ConditionalMongeGapEstimator: - r"""Monge Gap Estimator which optimizes over multiple conditions - - .. math:: - \text{min}_{\theta}\; \sum_{i=1}^{K} \Delta(T_\theta \sharp \mu, \theta) - + \lambda R(T_\theta \sharp \rho, \rho) - - Args: - dim_data: input dimensionality of data required for network init. - dim_cond: input dimensionality of condition embedding, - required for network init - model: network architecture for map :math:`T`, - should be a `ConditionalPerturbationNetwork`. - optimizer: optimizer function for map :math:`T`. - fitting_loss: function that outputs a fitting loss :math:`\Delta` between - two families of points, as well as any log object. - regularizer: function that outputs a score from two families of points, - here assumed to be of the same size, as well as any log object. - regularizer_strength: strength of the :attr:`regularizer`. - num_train_iters: number of total training iterations. - logging: option to return logs. - valid_freq: frequency with training and validation are logged. - rng: random key used for seeding for network initializations. - """ - - def __init__( - self, - dim_data: int, - dim_cond: int, - model: ConditionalPerturbationNetwork, - optimizer: Optional[optax.OptState] = None, - fitting_loss: Optional[ - Callable[[jnp.ndarray, jnp.ndarray], Tuple[float, Optional[Any]]] - ] = None, - regularizer: Optional[ - Callable[[jnp.ndarray, jnp.ndarray], Tuple[float, Optional[Any]]] - ] = None, - regularizer_strength: Union[float, Sequence[float]] = 1.0, - num_train_iters: int = 10_000, - logging: bool = False, - valid_freq: int = 500, - grad_acc_steps: int = 1, - rng: Optional[jax.Array] = None, - ) -> None: - self._fitting_loss = fitting_loss - self._regularizer = regularizer - self.num_train_iters = self.num_train_iters - self.grad_acc_steps = grad_acc_steps - # Can use either a fixed strength, or generalize to a schedule. - self.regularizer_strength = jnp.repeat( - jnp.atleast_2d(regularizer_strength), - num_train_iters, - total_repeat_length=num_train_iters, - axis=0, - ).ravel() - self.logging = logging - self.valid_freq = valid_freq - self.rng = jax.random.PRNGKey(rng) - - # set default optimizer - if optimizer is None: - optimizer = optax.adam(learning_rate=0.001) - - # setup training - self.setup(dim_data, dim_cond, model, optimizer) - - def setup( - self, - dim_data: int, - dim_cond: int, - neural_net: ConditionalPerturbationNetwork, - optimizer: optax.OptState, - ): - """Setup all components required to train the network""" + for c in jnp.unique(condition): + c_target = target[condition == c] + c_source = source[condition == c] - # neural network - self.rng, rng = jax.random.split(self.key, 2) - self.state_neural_net = neural_net.create_train_state( - self.rng, optimizer, dim_data, dim_cond - ) - - # step function - self.step_fn = self._get_step_fn() - - @property - def regularizer(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: - """Regularizer added to the fitting loss. - - Can be, e.g. the - :func:`~ott.neural.methods.monge_gap.monge_gap_from_samples`. - If no regularizer is passed for solver instantiation, - or regularization weight :attr:`regularizer_strength` is 0, - return 0 by default along with an empty set of log values. - """ # noqa: E501 - if self._regularizer is not None: - return self._regularizer - return lambda *_, **__: (0.0, None) - - @property - def fitting_loss(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: - """Fitting loss to fit the marginal constraint. - - Can be, e.g. :func:`~ott.tools.sinkhorn_divergence.sinkdiv`. - If no fitting_loss is passed for solver instantiation, return 0 by default, - and no log values. - """ - if self._fitting_loss is not None: - return self._fitting_loss - return lambda *_, **__: (0.0, None) - - def _generate_batch( - self, loader_source, loader_target - ) -> Dict[str, jnp.ndarray]: - """Generate a batch of condition and samples.""" - return { - "source": next(loader_source), - "target": next(loader_target), - } - - def train_map_estimator( - self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - ): - """The dataloaders should return a dict with key `X`. - The target dataloaders should additionally include a key - `c_embed`, which has the embedded condition in `dim_cond`.""" - # define logs - logs = collections.defaultdict(lambda: collections.defaultdict(list)) - - # try to display training progress with tqdm - try: - from tqdm import trange - - tbar = trange(self.num_train_iters, leave=True) - except ImportError: - tbar = range(self.num_train_iters) - - grads = tree_map(jnp.zeros_like, self.state_neural_net.params) - for step in tbar: - # update step - is_logging_step = self.logging and ( - (step % self.valid_freq == 0) - or (step == self.num_train_iters - 1) - ) - is_gradient_acc_step = (step + 1) % self.grad_acc_steps == 0 - train_batch, condition = self._generate_batch( - trainloader_source, trainloader_target - ) - valid_batch, _ = ( - None - if not is_logging_step - else self._generate_batch( - validloader_source, validloader_target - ) - ) - - self.state_neural_net, grads, current_logs = self.step_fn( - self.state_neural_net, - grads=grads, - train_batch=train_batch, - valid_batch=valid_batch, - is_logging_step=is_logging_step, - is_gradient_acc_step=is_gradient_acc_step, - ) - - # store and print metrics if logging step - if is_logging_step: - for log_key in current_logs: - for metric_key in current_logs[log_key]: - logs[log_key][metric_key].append( - current_logs[log_key][metric_key] - ) - - # update the tqdm bar if tqdm is available - if not isinstance(tbar, range): - reg_msg = ( - "NA" - if current_logs["eval"]["regularizer"] == 0.0 - else f"{current_logs['eval']['regularizer']:.4f}" - ) - postfix_str = ( - f"fitting_loss: {current_logs['eval']['fitting_loss']:.4f}, " - f"regularizer: {reg_msg} ," - f"total: {current_logs['eval']['total_loss']:.4f}" - ) - tbar.set_postfix_str(postfix_str) - - return self.state_neural_net, logs - - def _get_step_fn(self) -> Callable: - """Create a one step training and evaluation function.""" - - def loss_fn( - params: frozen_dict.FrozenDict, - apply_fn: Callable, - batch: Dict[str, jnp.ndarray], - ) -> Tuple[float, Dict[str, float]]: - """Loss function.""" - # map samples with the fitted map - mapped_samples = apply_fn( - {"params": params}, - batch["source"]["X"], - batch["target"]["c_embed"], - ) - - # compute the loss - val_fitting_loss, log_fitting_loss = self.fitting_loss( - batch["target"]["X"], mapped_samples - ) - val_regularizer, log_regularizer = self.regularizer( - batch["source"]["X"], mapped_samples - ) - val_tot_loss, log_regularizer = val_fitting_loss + val_regularizer - - # store training logs - loss_logs = { - "total_loss": val_tot_loss, - "fitting_loss": val_fitting_loss, - "regularizer": val_regularizer, - "log_regularizer": log_regularizer, - "log_fitting": log_fitting_loss, - } - - return val_tot_loss, loss_logs - - @functools.partial(jax.jit, static_argnums=[4, 5]) - def step_fn( - state_neural_net: train_state.TrainState, - grads: frozen_dict.FrozenDict, - train_batch: Dict[str, jnp.ndarray], - valid_batch: Optional[Dict[str, jnp.ndarray]] = None, - is_logging_step: bool = False, - is_gradient_acc_step: bool = False, - ) -> Tuple[ - train_state.TrainState, frozen_dict.FrozenDict, Dict[str, float] - ]: - """Step function.""" - # compute loss and gradients - grad_fn = jax.value_and_grad(loss_fn, argnums=0, has_aux=True) - (_, current_train_logs), step_grads = grad_fn( - state_neural_net.params, - state_neural_net.apply_fn, - train_batch, - ) - # Accumulate gradients - grads = tree_map(lambda g, step_g: g + step_g, grads, step_grads) - - # logging step - current_logs = {"train": current_train_logs, "eval": {}} - if is_logging_step: - _, current_eval_logs = loss_fn( - params=state_neural_net.params, - apply_fn=state_neural_net.apply_fn, - batch=valid_batch, - ) - current_logs["eval"] = current_eval_logs - - # update state - if is_gradient_acc_step: - state_neural_net = state_neural_net.apply_gradients( - grads=tree_map(lambda g: g / self.grad_acc_steps, grads) - ) - # Reset gradients - grads = tree_map(jnp.zeros_like, grads) - - return state_neural_net, grads, current_logs - - return step_fn - - def transport(self, x, c): - return self.state_neural_net.apply_fn( - {"params": self.state_neural_net.params}, x, c - ) - - @property - def model(self) -> nn.Module: - return self.state_neural_net - - @model.setter - def model(self, value: nn.Module): - """Setter for the model to be checkpointed.""" - self.state_neural_net = value - - def save_checkpoint(self, path: Optional[Path] = None) -> None: - """Abstract method for saving model parameters to a pickle file. - - Args: - path: Path where the checkpoint should be saved. Defaults to None in which case - it is retrieved from config. - config: The model training configuration with a `checkpointing_path` field. - Defaults to None. - NOTE: If `config` and `path` are both not None, `path` takes preference. - """ - if path is None: - raise ValueError( - "Checkpoint cannot be saved. Provide a checkpoint save path" - ) - try: - checkpointer = PyTreeCheckpointer() - save_args = save_args_from_target(self.model) - checkpointer.save(path, self.model, save_args=save_args, force=True) - except Exception as e: - raise Exception(f"Error in saving checkpoint to {path}: {e}") - - @classmethod - def load_checkpoint( - cls: Type[T], - ckpt_path: Path = None, - *args, - **kwargs, - ) -> T: - """ - Loading a model from a checkpoint - - Args: - cls: Class object to be created. - ckpt_path: Optional path from where checkpoint is restored. - Defaults to None, in that case inferred from config. - *args: args normally given to `ConditionalMongeGapEstimator` - - Returns: - Class object with restored weights. - """ - try: - out_class = cls( - *args, - **kwargs, - ) - checkpointer = PyTreeCheckpointer() - out_class.model = checkpointer.restore( - ckpt_path, item=out_class.model + if return_output: + monge_gap, out = monge_gap_from_samples( + target=c_target, source=c_source, return_output=True, **kwargs ) - return out_class - except Exception as e: - raise Exception( - f"Failed to load checkpoin from {ckpt_path}: {e}\nAre you sure" - "checkpoint was saved and correct path is provided?" + all_outs.append(out) + else: + monge_gap = monge_gap_from_samples( + target=c_target, source=c_source, return_output=False, **kwargs ) + all_gaps.append(monge_gap) + condition_monge_gap = sum(all_gaps) / len(all_gaps) # average -# Add condition embedding to dataloader --> Users -# How to sample condition...? -# Optim/embedding/regularization factory? + return ( + (condition_monge_gap, all_outs, all_gaps) + if return_output + else condition_monge_gap + ) From ebdaaebb4fb777c1ed295d96f4098a5cc62964ae Mon Sep 17 00:00:00 2001 From: Alice Driessen Date: Mon, 5 May 2025 13:23:04 +0200 Subject: [PATCH 4/6] Working with some inplace hacks - not pretty --- src/ott/neural/methods/monge_gap.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/ott/neural/methods/monge_gap.py b/src/ott/neural/methods/monge_gap.py index 31da10f55..ddaceb9b6 100644 --- a/src/ott/neural/methods/monge_gap.py +++ b/src/ott/neural/methods/monge_gap.py @@ -219,6 +219,7 @@ def __init__( logging: bool = False, valid_freq: int = 500, rng: Optional[jax.Array] = None, + dim_cond: Optional[int]=None, ): self._fitting_loss = fitting_loss self._regularizer = regularizer @@ -239,18 +240,19 @@ def __init__( optimizer = optax.adam(learning_rate=0.001) # setup training - self.setup(dim_data, model, optimizer) + self.setup(dim_data, model, optimizer, dim_cond) def setup( self, dim_data: int, neural_net: potentials.BasePotential, optimizer: optax.OptState, + dim_cond: Union[None, int] ): """Setup all components required to train the network.""" # neural network self.state_neural_net = neural_net.create_train_state( - self.rng, optimizer, dim_data + self.rng, optimizer, dim_data, dim_cond ) # step function @@ -365,14 +367,14 @@ def loss_fn( ) -> Tuple[float, Dict[str, float]]: """Loss function.""" # map samples with the fitted map - mapped_samples = apply_fn({"params": params}, batch["source"]) + mapped_samples = apply_fn({"params": params}, batch["source"]["X"], batch["source"]["c"]) # compute the loss val_fitting_loss, log_fitting_loss = self.fitting_loss( mapped_samples, batch["target"] ) val_regularizer, log_regularizer = self.regularizer( - batch["source"], mapped_samples + batch["source"]["X"], mapped_samples, batch["target"]["c"] ) val_tot_loss = ( val_fitting_loss + self.regularizer_strength[step] * val_regularizer From 4f68fbd57315616a5d5ceb830864037688fccf12 Mon Sep 17 00:00:00 2001 From: Alice Driessen Date: Thu, 15 May 2025 11:24:44 +0200 Subject: [PATCH 5/6] No changes to Monge Gap --- src/ott/neural/methods/monge_gap.py | 12 +++++----- .../conditional_perturbation_network.py | 22 ++++++++++++++----- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/src/ott/neural/methods/monge_gap.py b/src/ott/neural/methods/monge_gap.py index ddaceb9b6..afe8dc273 100644 --- a/src/ott/neural/methods/monge_gap.py +++ b/src/ott/neural/methods/monge_gap.py @@ -219,7 +219,7 @@ def __init__( logging: bool = False, valid_freq: int = 500, rng: Optional[jax.Array] = None, - dim_cond: Optional[int]=None, + # dim_cond: Optional[int]=None, ): self._fitting_loss = fitting_loss self._regularizer = regularizer @@ -240,19 +240,19 @@ def __init__( optimizer = optax.adam(learning_rate=0.001) # setup training - self.setup(dim_data, model, optimizer, dim_cond) + self.setup(dim_data, model, optimizer) def setup( self, dim_data: int, neural_net: potentials.BasePotential, optimizer: optax.OptState, - dim_cond: Union[None, int] + # dim_cond: Union[None, int] ): """Setup all components required to train the network.""" # neural network self.state_neural_net = neural_net.create_train_state( - self.rng, optimizer, dim_data, dim_cond + self.rng, optimizer, dim_data ) # step function @@ -367,14 +367,14 @@ def loss_fn( ) -> Tuple[float, Dict[str, float]]: """Loss function.""" # map samples with the fitted map - mapped_samples = apply_fn({"params": params}, batch["source"]["X"], batch["source"]["c"]) + mapped_samples = apply_fn({"params": params}, batch["source"]) # compute the loss val_fitting_loss, log_fitting_loss = self.fitting_loss( mapped_samples, batch["target"] ) val_regularizer, log_regularizer = self.regularizer( - batch["source"]["X"], mapped_samples, batch["target"]["c"] + batch["source"], mapped_samples ) val_tot_loss = ( val_fitting_loss + self.regularizer_strength[step] * val_regularizer diff --git a/src/ott/neural/networks/conditional_perturbation_network.py b/src/ott/neural/networks/conditional_perturbation_network.py index 81fefff76..9b63ce6b8 100644 --- a/src/ott/neural/networks/conditional_perturbation_network.py +++ b/src/ott/neural/networks/conditional_perturbation_network.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Iterable, Sequence, Tuple +from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple, Union import flax.linen as nn import jax.numpy as jnp @@ -30,8 +30,8 @@ class ConditionalPerturbationNetwork(BasePotential): @nn.compact def __call__( - self, x: jnp.ndarray, c: jnp.ndarray - ) -> jnp.ndarray: # noqa: D102 + self, x: jnp.ndarray, c: Optional[jnp.ndarray]=None + ) -> Union[jnp.ndarray, Dict[str, jnp.ndarray]]: # noqa: D102 """ Args: x (jnp.ndarray): The input data of shape bs x dim_data @@ -42,6 +42,12 @@ def __call__( Returns: jnp.ndarray: _description_ """ + return_batch = False + if isinstance(x, dict): + c = x["c"] + x = x["X"] + return_batch = True + n_input = x.shape[-1] # Chunk the inputs @@ -93,18 +99,22 @@ def __call__( z = self.act_fn(wx(z)) wx = nn.Dense(n_input, use_bias=True) - return x + wx(z) + y = x + wx(z) + + if return_batch: + return {"X": y, "c": c} + else: + return y def create_train_state( self, rng: jnp.ndarray, optimizer: optax.OptState, dim_data: int, - dim_cond: int, **kwargs: Any, ) -> PotentialTrainState: """Create initial `TrainState`.""" - c = jnp.ones((1, dim_cond)) # (n_batch, embed_dim) + c = jnp.ones((1, self.dim_cond)) # (n_batch, embed_dim) x = jnp.ones((1, dim_data)) # (n_batch, data_dim) params = self.init(rng, x=x, c=c)["params"] return PotentialTrainState.create( From e2670b2d826615b627361c4081c7d3a6082a46e5 Mon Sep 17 00:00:00 2001 From: Alice Driessen Date: Thu, 15 May 2025 11:26:10 +0200 Subject: [PATCH 6/6] Removed commented code --- src/ott/neural/methods/monge_gap.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/ott/neural/methods/monge_gap.py b/src/ott/neural/methods/monge_gap.py index afe8dc273..31da10f55 100644 --- a/src/ott/neural/methods/monge_gap.py +++ b/src/ott/neural/methods/monge_gap.py @@ -219,7 +219,6 @@ def __init__( logging: bool = False, valid_freq: int = 500, rng: Optional[jax.Array] = None, - # dim_cond: Optional[int]=None, ): self._fitting_loss = fitting_loss self._regularizer = regularizer @@ -247,7 +246,6 @@ def setup( dim_data: int, neural_net: potentials.BasePotential, optimizer: optax.OptState, - # dim_cond: Union[None, int] ): """Setup all components required to train the network.""" # neural network