指标

class flax.nnx.metrics.Metric(*args, **kwargs)

指标的基类。任何子类化 Metric 的类都应实现 compute、reset 和 update 方法。

__init__()

compute()

计算并返回 Metric 的值。

reset()

就地重置 Metric。

update(**kwargs)

就地更新 Metric。

class flax.nnx.metrics.Average(*args, **kwargs)

平均指标。

示例用法

>>> import jax.numpy as jnp
>>> from flax import nnx

>>> batch_loss = jnp.array([1, 2, 3, 4])
>>> batch_loss2 = jnp.array([3, 2, 1, 0])

>>> metrics = nnx.metrics.Average()
>>> metrics.compute()
Array(nan, dtype=float32)
>>> metrics.update(values=batch_loss)
>>> metrics.compute()
Array(2.5, dtype=float32)
>>> metrics.update(values=batch_loss2)
>>> metrics.compute()
Array(2., dtype=float32)
>>> metrics.reset()
>>> metrics.compute()
Array(nan, dtype=float32)

__init__(argname='values')

传入一个字符串，表示 update() 将用于推断新值的关键字参数。例如，将度量标准构造为 avg = Average('test') 将允许您使用 avg.update(test=new_value) 进行更新。

参数

argname – 一个可选字符串，表示 update() 将用于推断新值的关键字参数。默认为 'values'。

compute()

计算并返回平均值。

reset()

重置此 Metric。

update(**kwargs)

就地更新此 Metric。此方法将使用来自 kwargs[self.argname] 的值来更新度量标准，其中 self.argname 在构造时定义。

参数

**kwargs – 包含 self.argname 条目的关键字参数，它映射到我们要用来更新此度量标准的值。

class flax.nnx.metrics.Accuracy(*args, **kwargs)

准确率指标。此指标是 Average 的子类，因此它们共享相同的 reset 和 compute 方法实现。与 Average 不同，构造时不需要向 Accuracy 传入字符串。

示例用法

>>> from flax import nnx
>>> import jax, jax.numpy as jnp

>>> logits = jax.random.normal(jax.random.key(0), (5, 2))
>>> labels = jnp.array([1, 1, 0, 1, 0])
>>> logits2 = jax.random.normal(jax.random.key(1), (5, 2))
>>> labels2 = jnp.array([0, 1, 1, 1, 1])

>>> metrics = nnx.metrics.Accuracy()
>>> metrics.compute()
Array(nan, dtype=float32)
>>> metrics.update(logits=logits, labels=labels)
>>> metrics.compute()
Array(0.6, dtype=float32)
>>> metrics.update(logits=logits2, labels=labels2)
>>> metrics.compute()
Array(0.7, dtype=float32)
>>> metrics.reset()
>>> metrics.compute()
Array(nan, dtype=float32)

update(*, logits, labels, **_)

就地更新此 Metric。

参数

• logits – 输出的预测激活。在将这些值与标签进行比较之前，会对其（在尾随维度上）进行 argmax 操作。

• labels – 真实整数标签。

class flax.nnx.metrics.Welford(*args, **kwargs)

使用 Welford 算法来计算数据流的均值和方差。

示例用法

>>> import jax.numpy as jnp
>>> from flax import nnx

>>> batch_loss = jnp.array([1, 2, 3, 4])
>>> batch_loss2 = jnp.array([3, 2, 1, 0])

>>> metrics = nnx.metrics.Welford()
>>> metrics.compute()
Statistics(mean=Array(0., dtype=float32), standard_error_of_mean=Array(nan, dtype=float32), standard_deviation=Array(nan, dtype=float32))
>>> metrics.update(values=batch_loss)
>>> metrics.compute()
Statistics(mean=Array(2.5, dtype=float32), standard_error_of_mean=Array(0.559017, dtype=float32), standard_deviation=Array(1.118034, dtype=float32))
>>> metrics.update(values=batch_loss2)
>>> metrics.compute()
Statistics(mean=Array(2., dtype=float32), standard_error_of_mean=Array(0.43301272, dtype=float32), standard_deviation=Array(1.2247449, dtype=float32))
>>> metrics.reset()
>>> metrics.compute()
Statistics(mean=Array(0., dtype=float32), standard_error_of_mean=Array(nan, dtype=float32), standard_deviation=Array(nan, dtype=float32))

__init__(argname='values')

传入一个字符串，表示 update() 将用于推断新值的关键字参数。例如，将度量标准构造为 wf = Welford('test') 将允许您使用 wf.update(test=new_value) 进行更新。

参数

argname – 一个可选字符串，表示 update() 将用于推断新值的关键字参数。默认为 'values'。

compute()

计算并返回 Statistics 数据类对象中的均值和方差统计信息。

reset()

重置此 Metric。

update(**kwargs)

就地更新此 Metric。此方法将使用来自 kwargs[self.argname] 的值来更新度量标准，其中 self.argname 在构造时定义。

参数

**kwargs – 包含 self.argname 条目的关键字参数，它映射到我们要用来更新此度量标准的值。

class flax.nnx.metrics.MultiMetric(*args, **kwargs)

MultiMetric 类用于存储多个指标并在单个调用中更新它们。

示例用法

>>> from flax import nnx
>>> import jax, jax.numpy as jnp

>>> metrics = nnx.MultiMetric(
...   accuracy=nnx.metrics.Accuracy(), loss=nnx.metrics.Average()
... )

>>> metrics
MultiMetric(
  accuracy=Accuracy(
    argname='values',
    total=MetricState(
      value=Array(0., dtype=float32)
    ),
    count=MetricState(
      value=Array(0, dtype=int32)
    )
  ),
  loss=Average(
    argname='values',
    total=MetricState(
      value=Array(0., dtype=float32)
    ),
    count=MetricState(
      value=Array(0, dtype=int32)
    )
  )
)

>>> metrics.accuracy
Accuracy(
  argname='values',
  total=MetricState(
    value=Array(0., dtype=float32)
  ),
  count=MetricState(
    value=Array(0, dtype=int32)
  )
)

>>> metrics.loss
Average(
  argname='values',
  total=MetricState(
    value=Array(0., dtype=float32)
  ),
  count=MetricState(
    value=Array(0, dtype=int32)
  )
)

>>> logits = jax.random.normal(jax.random.key(0), (5, 2))
>>> labels = jnp.array([1, 1, 0, 1, 0])
>>> logits2 = jax.random.normal(jax.random.key(1), (5, 2))
>>> labels2 = jnp.array([0, 1, 1, 1, 1])

>>> batch_loss = jnp.array([1, 2, 3, 4])
>>> batch_loss2 = jnp.array([3, 2, 1, 0])

>>> metrics.compute()
{'accuracy': Array(nan, dtype=float32), 'loss': Array(nan, dtype=float32)}
>>> metrics.update(logits=logits, labels=labels, values=batch_loss)
>>> metrics.compute()
{'accuracy': Array(0.6, dtype=float32), 'loss': Array(2.5, dtype=float32)}
>>> metrics.update(logits=logits2, labels=labels2, values=batch_loss2)
>>> metrics.compute()
{'accuracy': Array(0.7, dtype=float32), 'loss': Array(2., dtype=float32)}
>>> metrics.reset()
>>> metrics.compute()
{'accuracy': Array(nan, dtype=float32), 'loss': Array(nan, dtype=float32)}

__init__(**metrics)

向构造函数传入关键字参数，例如 MultiMetric(keyword1=Average(), keyword2=Accuracy(), ...)。

参数

**metrics** – 将用于访问相应 Metric 的关键字参数。

compute()

计算并返回所有基础 Metric 的值。此方法将返回一个字典，将字符串（由传递给构造函数的关键字参数 **metrics 定义）映射到相应的度量值。

reset()

重置所有基础 Metric。

update(**updates)

就地更新此 MultiMetric 中的所有基础 Metric。所有 **updates 将传递给所有基础 Metric 的 update 方法。

参数

**updates – 将传递给基础 Metric 的 update 方法的关键字参数。