TorchMetrics

TorchMetrics is a collection of 80+ PyTorch metrics implementations and an easy-to-use API to create custom metrics. Your data will always be placed on the same device as your metrics. You can log Metric objects directly in Lightning to reduce even more boilerplate. The module-based metrics contain internal metric states (similar to the parameters of the PyTorch module) that automate accumulation and synchronization across devices! Automatic accumulation over multiple batches. Automatic synchronization between multiple devices. Metric arithmetic. Similar to torch.nn, most metrics have both a module-based and a functional version. The functional versions are simple python functions that as input take torch.tensors and return the corresponding metric as a torch.tensor.