We provide Cerebras-compatible metrics that can be used during evaluation to measure how well the model has been trained. These metrics can be found in the metrics.Metric module. For example: 
model = Model()
compiled_model = cstorch.compile(model, backend)

accuracy = cstorch.metrics.AccuracyMetric("accuracy")

@cstorch.trace
def eval_step(batch):
    inputs, targets = batch
    outputs = compiled_model(inputs)

    accuracy(
        labels=targets.clone(),
        predictions=outputs.argmax(-1).int(),
    )

...
for batch in executor:
    ...


# Log accumulated eval metric
print(f"Accuracy: {float(accuracy)}")

​
Writing Custom Metrics

To define a Cerebras compliant metrics, create a subclass of cerebras.pytorch.metrics.Metric. For example, 
class CustomMetric(cstorch.metrics.Metric):

    def __init__(self, name: str):
        super().__init__(name)
        ...

    ...

    def reset(self):
        pass

    def update(self, *args, **kwargs):
        ...

    def compute(self):
        ...
As can be seen in the above example, the base Metric class expects one argument. Namely, the metric name. In addition, there are three abstract methods that must be overridden:
  • reset This method resets (or defines if its the first time its called) the metrics’ internal state.States can be registered via calls to register_state
  • update This method is used to update the metric’s registered states.Note that to remain Cerebras compliant, no tensor may be evaluated/inspected here. The update call is intended to be fully traced.
  • compute This method is used to compute the final accumulated metric value using the state that was updated in update