sf.training.trainer.map_trainer¶

mrmustard.training.trainer.map_trainer(trainer=<function train_device>, tasks=1, pbar=True, unblock=False, num_cpus=None, **kwargs)[source]¶

Maps multiple training tasks across multiple workers using ray.

In practice, the most common use case is to ignore the keywords trainer (as it defaults to train_device()), pbar, unblock, etc. and just concentrate on tasks and **kwargs which passes arguments to the wrapper functions that contain the task execution logic, as well as the Optimizer and its Optimizer.minimize().

For example, with the default trainer train_device(), two user-defined functions are used for wrapping up user logic:

A device_factory (optional) that wraps around the logic for making circuits/states to be optimized; it is expected to return a single, or list of, :class:`Circuit`(s).
A cost_fn (required) that takes the circuits made and additional keyword arguments and returns a backprop-able scalar cost.

Refer to the kwargs section below for more available options.

Parameters:

trainer (callable) – The function containing the training loop to be distributed, whose fixed arguments are to be passed by **kwargs and task-specific arguments iterated through tasks. Provide only when custom evaluation/training logic is needed. Defaults to train_device().
tasks (Union[int, Sequence, Mapping]) – Number of repeats or collection of task-specific training config arguments feeding into train_device(). Refer to kwargs below for the available options. Defaults to 1 which runs trainer exactly once.
pbar (bool) – Whether to show a progress bar, available only in blocking mode (i.e. unblock==False). Defaults to True.
unblock (bool) – Whether to unblock the process and returns a getter function returning the available results. Defaults to False.
num_cpus (int) – Number of cpu workers to initialize ray. Defaults to the number of virtual cores.
kwargs –
Additional arguments containing fixed training config kwargs feeding into trainer. For the default trainer train_device(), available options are:
- cost_fn (callable):
  The optimized cost function to be distributed. It’s expected to accept the output of device_factory as *args as well as user-defined **kwargs, and returns a scalar cost. Its user-defined **kwargs will be passed from this function’s **kwargs which must include all its required arguments.
- device_factory (callable):
  Function that (partially) takes kwargs and returns a device, or list/dict of devices. If None, cost_fn will be assumed to take no positional argument (for example, when device-making is contained in cost_fn). Defaults to None.
- metric_fns (Union[Sequence[callable], Mapping[callable], callable]):
  Optional collection of functions that takes the output of device_factory after optimization and returns arbitrary evaluation/information.
- return_kwargs (bool):
  Whether to include input config kwargs in the output dict. Defualts to True.
- skip_opt (bool):
  Whether to skip the optimization and directly calculate cost.
- tag (str):
  Optional label of the training task associated with the kwargs to be included in the output dict.
- any kwargs to cost_fn: exluding the output of device_factory.
- any kwargs to device_factory: e.g. x, r, theta, etc.
- any kwargs to Optimizer: e.g. euclidean_lr.
- any kwargs to Optimizer.minimize: excluding cost_fn and by_optimizing, e.g. max_steps.

Returns

Union[List, Dict]: The collection of results from each training task. Returns

a list if tasks is provided as an int or a list; or
a dict with the same keys if tasks is provided as a dict.

Examples:¶

from mrmustard.lab import Vacuum, Dgate, Ggate, Gaussian
from mrmustard.physics import fidelity
from mrmustard.training.trainer import map_trainer

def make_circ(x=0.):
    return Ggate(num_modes=1, symplectic_trainable=True) >> Dgate(x=x, x_trainable=True, y_trainable=True)

def cost_fn(circ=make_circ(0.1), y_targ=0.):
    target = Gaussian(1) >> Dgate(-1.5, y_targ)
    s = Vacuum(1) >> circ
    return -fidelity(s, target)

# Use case 0: Calculate the cost of a randomly initialized circuit 5 times without optimizing it.
results_0 = map_trainer(
    cost_fn=cost_fn,
    tasks=5,
)

# Use case 1: Run circuit optimization 5 times on randomly initialized circuits.
results_1 = map_trainer(
    cost_fn=cost_fn,
    device_factory=make_circ,
    tasks=5,
    max_steps=50,
    symplectic_lr=0.05,
)

# Use case 2: Run 2 sets of circuit optimization with custom parameters passed as list.
results_2 = map_trainer(
    cost_fn=cost_fn,
    device_factory=make_circ,
    tasks=[
        {'x': 0.1, 'euclidean_lr': 0.005, 'max_steps': 50},
        {'x': -0.7, 'euclidean_lr': 0.1, 'max_steps': 2},
    ],
    y_targ=0.35,
    symplectic_lr=0.05,
    AUTOCUTOFF_MAX_CUTOFF=7,
)

# Use case 3: Run 2 sets of circuit optimization with custom parameters passed as dict with extra metric functions for evaluating the final optimized circuit.
results_3 = map_trainer(
cost_fn=cost_fn,
device_factory=make_circ,
tasks={
    'my-job': {'x': 0.1, 'euclidean_lr': 0.005, 'max_steps': 50},
    'my-other-job': {'x': -0.7, 'euclidean_lr': 0.1, 'max_steps': 2},
},
y_targ=0.35,
symplectic_lr=0.05,
metric_fns={
    'is_gaussian': lambda c: c.is_gaussian,
    'foo': lambda _: 17.
},

sf.training.trainer.map_trainer¶

Examples:¶

Contents