pgcuts.utils

Gradient mixing

class pgcuts.optim.GradientMixer(named_parameters, loss_scale)

Mix gradients from multiple losses.

Parameters:

• named_parameters (Iterator[Tuple[str, Parameter]])

• loss_scale (Dict[str, float])

__init__(named_parameters, loss_scale)

Initialize gradient mixer.

Parameters:

• named_parameters (Iterator[Tuple[str, Parameter]]) – Model named parameters.

• loss_scale (Dict[str, float]) – Dict of loss name to scale.

attach()

Attach gradient scaling hooks.

detach()

Remove all gradient hooks.

static make_grad_scaler_hook(scales_book)

Create gradient scaling hook.

Parameters:

scales_book (Dict[str, float]) – Dict with ‘loss_scale’ key.

Returns:

Hook function.

Usage:

from pgcuts.optim import GradientMixer

grad_mix = GradientMixer(
    network.named_parameters(),
    loss_scale={"cut": 1.0, "balance": 1.0},
)

optimizer.zero_grad()
with grad_mix("cut"):
    cut_loss.backward(retain_graph=True)
with grad_mix("balance"):
    balance_loss.backward()
optimizer.step()

Pair utilities

pgcuts.utils.pairs.get_pairs_unique_map(pairs)

Return unique indices and inverse maps.

Parameters:

pairs (ndarray) – Array of (i,j) pairs, shape (k, 2).

Returns:

Tuple of (unique, left_inv, right_inv).

Return type:

Tuple[ndarray, …]