Package: jax-metrics
JAX MetricsA Metrics library for the JAX ecosystemMain FeaturesStandard metrics that can be used in any JAX project.Pytree abstractions that can natively integrate with all JAX APIs and pytree-supporting frameworks (flax.struct, equinox, treex, etc).Distributed-friendly APIs that make it super easy to synchronize metrics across devices.Automatic accumulation over epochs.JAX Metrics is implemented on top ofTreeo.What is included?The Keras-likeLossandMetricabstractions.Ametricsmodule containing popular metrics.Thelossesandregularizersmodules containing popular losses.TheMetricsandLossescombinators.InstallationInstall using pip:pipinstalljax_metricsStatusMetrics on this library are usually tested against their Keras or Torchmetrics counterparts for numerical equivalence. This code base comes from Treex and Elegy so it's already in use.Getting StartedMetricTheMetricAPI consists of 3 basic methods:reset: Used to both initialize and reset a metric.update: Takes in new data and updates the metric state.compute: Returns the current value of the metric.Simple usage looks like this:importjax_metricsasjmmetric=jm.metrics.Accuracy()# Update the metric with a batch of predictions and labelsmetric=metric.update(target=y,preds=logits)# Get the current value of the metricacc=metric.compute()# 0.95# alternatively, produce a logs dictlogs=metric.compute_logs()# {'accuracy': 0.95}# Reset the metricmetric=metric.reset()Note thatupdateenforces the use of keyword arguments. Also theMetric.nameproperty is used as the key in the returned dict, by default this is the name of the class in lowercase but can be overridden in the constructor via thenameargument.Tipical Training SetupBecause Metrics are pytrees they can be used withjit,pmap, etc. On a more realistic scenario you will proably want to use them inside some of your JAX functions in a setup similar to this:importjax_metricsasjmmetric=jm.metrics.Accuracy()defloss_fn(params,metric,x,y):...metric=metric.update(target=y,preds=logits)...returnloss,metric@jax.jitdeftrain_step(params,metric,x,y):grads,metric=jax.grad(loss_fn,has_aux=True)(params,metric,x,y)...returnparams,metric@jax.jitdefreset_step(metric:jm.Metric)->jm.Metric:returnmetric.reset()Since the loss function usually has access to the predictions and labels, its usually where you would callmetric.update, and the new metric state can be returned as an auxiliary output.Distributed TrainingJAX Metrics has a distributed friendly API via thebatch_updatesandreducemethods. A simple example of a loss function inside a data parallel setup could look like this:defloss_fn(params,metric,x,y):...# compuate batch updatebatch_updates=metric.batch_updates(target=y,preds=logits)# gather over all devices and reducebatch_updates=jax.lax.all_gather(batch_updates,"device").reduce()# update metricmetric=metric.merge(batch_updates)...Thebatch_updatesmethod behaves similar toupdatebut returns a new metric state with only information about that batch,jax.lax.all_gather"gathers" the metric state over all devices plus adds a new axis to the metric state, andreducereduces the metric state over all devices (first axis). Finally,mergecombines the accumulated metric state over the previous batches with the batch updates.LossTheLossAPI just consists of a__call__method. Simple usage looks like this:importjax_metricsasjmcrossentropy=jm.losses.Crossentropy()# get reduced loss valueloss=crossentropy(target=y,preds=logits)# 0.23Note that losses are not pytrees so they should be marked as static. Similar to Keras, all losses have areductionstrategy that can be specified in the constructor and (usually) makes sure that the output is a scalar.Why have losses in a metrics library?There are a few reasons for having losses in a metrics library:Most code from this library was originally written for and will still be consumed by Elegy. Since Elegy needs support for calculating cumulative losses, as you will see later, a Metric abstraction calledLosseswas created for this.A couple of API design decisions are shared between theLossandMetricAPIs. This includes:__call__andupdateboth accept any number keyword only arguments. This is used to facilitate composition (seeCombinatorssection).Both classes have theindex_intoandrename_argumentsmethods that allow them to modify how arguments are consumed.Argument names are standardized to be consistent when ever possible, e.g. bothmetrics.Accuracyandlosses.Crossentropyuse thetargetandpredsarguments.CombinatorsCombinators enable you to group together multiple metrics while also being instances ofMetricand thus maintaining the same API.MetricsTheMetricscombinator lets you combine multiple metrics into a single Metric object.metrics=jm.Metrics([jm.metrics.Accuracy(),jm.metrics.F1(),# not yet implemented 😅, coming soon?])# same APImetrics=metrics.update(target=y,preds=logits)# compute now returns a dictmetrics.compute()# {'accuracy': 0.95, 'f1': 0.87}# same as compute_logs in the casemetrics.compute_logs()# {'accuracy': 0.95, 'f1': 0.87}# Reset the metricsmetrics=metrics.reset()As you can see theMetrics.updatemethod accepts and forwards all the arguments required by the individual metrics. In this example they use the same arguments, but in practice they may consume different subsets of the arguments. Also, if names are repeated then unique names are generated for each metric by appending a number to the metric name.If a dictionary is used instead of a list, the keys are used instead of thenameproperty of the metrics to determine the key in the returned dict.metrics=jm.Metrics({"acc":jm.metrics.Accuracy(),"f_one":jm.metrics.F1(),# not yet implemented 😅, coming soon?})# same APImetrics=metrics.update(target=y,preds=logits)# compute new returns a dictmetrics.compute()# {'acc': 0.95, 'f_one': 0.87}# same as compute_logs in the casemetrics.compute_logs()# {'acc': 0.95, 'f_one': 0.87}# Reset the metricsmetrics=metrics.reset()You can use nested structures of dicts and lists to group metrics, the keys of the dicts are used to determine group names. Group names and metrics names are concatenated using"/"e.g."{group_name}/{metric_name}".LossesLossesis aMetriccombinator that behaves very similarly toMetricsbut containsLossinstances.Lossescalculates the cumulativemeanvalue of each loss over the batches.losses=jm.Losses([jm.losses.Crossentropy(),jm.regularizers.L2(1e-4),])# same APIlosses=losses.update(target=y,preds=logits,parameters=params)# compute new returns a dictlosses.compute()# {'crossentropy': 0.23, 'l2': 0.005}# same as compute_logs in the caselosses.compute_logs()# {'crossentropy': 0.23, 'l2': 0.005}# you can also compute the total lossloss=losses.total_loss()# 0.235# Reset the losseslosses=losses.reset()As withMetrics, theupdatemethod accepts and forwards all the arguments required by the individual losses. In this exampletargetandpredsare used by theCrossentropy, whileparametersis used by theL2. Thetotal_lossmethod returns the sum of all values returned bycompute.If a dictionary is used instead of a list, the keys are used instead of thenameproperty of the losses to determine the key in the returned dict.losses=jm.Losses({"xent":jm.losses.Crossentropy(),"l2":jm.regularizers.L2(1e-4),})# same APIlosses=losses.update(target=y,preds=logits,parameters=params)# compute new returns a dictlosses.compute()# {'xent': 0.23, 'l2': 0.005}# same as compute_logs in the caselosses.compute_logs()# {'xent': 0.23, 'l2': 0.005}# you can also compute the total lossloss=losses.total_loss()# 0.235# Reset the losseslosses=losses.reset()If you want to useLossesto calculate the loss of a model, you should usebatch_updatesfollowed bytotal_lossto get the correct batch loss. For example, a loss function could be written as:defloss_fn(...,losses):...batch_updates=losses.batch_updates(target=y,preds=logits,parameters=params)loss=batch_updates.total_loss()losses=losses.merge(batch_updates)...returnloss,lossesFor convenience, the previous pattern can be simplified to a single line using theloss_and_updatemethod:defloss_fn(...):...loss,lossses=losses.loss_and_update(target=y,preds=logits,parameters=params)...returnloss,losses