Package: torchmetrics
Machine learning metrics for distributed, scalable PyTorch applications.What is Torchmetrics•Implementing a metric•Built-in metrics•Docs•Community•LicenseInstallationSimple installation from PyPIpipinstalltorchmetricsOther installationsInstall using condacondainstall-cconda-forgetorchmetricsPip from source# with gitpipinstallgit+https://github.com/Lightning-AI/torchmetrics.git@release/stablePip from archivepipinstallhttps://github.com/Lightning-AI/torchmetrics/archive/refs/heads/release/stable.zipExtra dependencies for specialized metrics:pipinstalltorchmetrics[audio]pipinstalltorchmetrics[image]pipinstalltorchmetrics[text]pipinstalltorchmetrics[all]# install all of the aboveInstall latest developer versionpipinstallhttps://github.com/Lightning-AI/torchmetrics/archive/master.zipWhat is TorchMetricsTorchMetrics is a collection of 100+ PyTorch metrics implementations and an easy-to-use API to create custom metrics. It offers:A standardized interface to increase reproducibilityReduces boilerplateAutomatic accumulation over batchesMetrics optimized for distributed-trainingAutomatic synchronization between multiple devicesYou can use TorchMetrics with any PyTorch model or withPyTorch Lightningto enjoy additional features such as:Module metrics are automatically placed on the correct device.Native support for logging metrics in Lightning to reduce even more boilerplate.Using TorchMetricsModule metricsThemodule-based metricscontain internal metric states (similar to the parameters of the PyTorch module) that automate accumulation and synchronization across devices!Automatic accumulation over multiple batchesAutomatic synchronization between multiple devicesMetric arithmeticThis can be run on CPU, single GPU or multi-GPUs!For the single GPU/CPU case:importtorch# import our libraryimporttorchmetrics# initialize metricmetric=torchmetrics.classification.Accuracy(task="multiclass",num_classes=5)# move the metric to device you want computations to take placedevice="cuda"iftorch.cuda.is_available()else"cpu"metric.to(device)n_batches=10foriinrange(n_batches):# simulate a classification problempreds=torch.randn(10,5).softmax(dim=-1).to(device)target=torch.randint(5,(10,)).to(device)# metric on current batchacc=metric(preds,target)print(f"Accuracy on batch{i}:{acc}")# metric on all batches using custom accumulationacc=metric.compute()print(f"Accuracy on all data:{acc}")Module metric usage remains the same when using multiple GPUs or multiple nodes.Example using DDPimportosimporttorchimporttorch.distributedasdistimporttorch.multiprocessingasmpfromtorchimportnnfromtorch.nn.parallelimportDistributedDataParallelasDDPimporttorchmetricsdefmetric_ddp(rank,world_size):os.environ["MASTER_ADDR"]="localhost"os.environ["MASTER_PORT"]="12355"# create default process groupdist.init_process_group("gloo",rank=rank,world_size=world_size)# initialize modelmetric=torchmetrics.classification.Accuracy(task="multiclass",num_classes=5)# define a model and append your metric to it# this allows metric states to be placed on correct accelerators when# .to(device) is called on the modelmodel=nn.Linear(10,10)model.metric=metricmodel=model.to(rank)# initialize DDPmodel=DDP(model,device_ids=[rank])n_epochs=5# this shows iteration over multiple training epochsforninrange(n_epochs):# this will be replaced by a DataLoader with a DistributedSamplern_batches=10foriinrange(n_batches):# simulate a classification problempreds=torch.randn(10,5).softmax(dim=-1)target=torch.randint(5,(10,))# metric on current batchacc=metric(preds,target)ifrank==0:# print only for rank 0print(f"Accuracy on batch{i}:{acc}")# metric on all batches and all accelerators using custom accumulation# accuracy is same across both acceleratorsacc=metric.compute()print(f"Accuracy on all data:{acc}, accelerator rank:{rank}")# Resetting internal state such that metric ready for new datametric.reset()# cleanupdist.destroy_process_group()if__name__=="__main__":world_size=2# number of gpus to parallelize overmp.spawn(metric_ddp,args=(world_size,),nprocs=world_size,join=True)Implementing your own Module metricImplementing your own metric is as easy as subclassing antorch.nn.Module. Simply, subclasstorchmetrics.Metricand just implement theupdateandcomputemethods:importtorchfromtorchmetricsimportMetricclassMyAccuracy(Metric):def__init__(self):# remember to call supersuper().__init__()# call `self.add_state`for every internal state that is needed for the metrics computations# dist_reduce_fx indicates the function that should be used to reduce# state from multiple processesself.add_state("correct",default=torch.tensor(0),dist_reduce_fx="sum")self.add_state("total",default=torch.tensor(0),dist_reduce_fx="sum")defupdate(self,preds:torch.Tensor,target:torch.Tensor)->None:# extract predicted class index for computing accuracypreds=preds.argmax(dim=-1)assertpreds.shape==target.shape# update metric statesself.correct+=torch.sum(preds==target)self.total+=target.numel()defcompute(self)->torch.Tensor:# compute final resultreturnself.correct.float()/self.totalmy_metric=MyAccuracy()preds=torch.randn(10,5).softmax(dim=-1)target=torch.randint(5,(10,))print(my_metric(preds,target))Functional metricsSimilar totorch.nn, most metrics have both amodule-basedand functional version.
The functional versions are simple python functions that as input taketorch.tensorsand return the corresponding metric as atorch.tensor.importtorch# import our libraryimporttorchmetrics# simulate a classification problempreds=torch.randn(10,5).softmax(dim=-1)target=torch.randint(5,(10,))acc=torchmetrics.functional.classification.multiclass_accuracy(preds,target,num_classes=5)Covered domains and example metricsIn total TorchMetrics contains100+ metrics, which
covers the following domains:AudioClassificationDetectionInformation RetrievalImageMultimodal (Image-Text)NominalRegressionTextEach domain may require some additional dependencies which can be installed withpip install torchmetrics[audio],pip install torchmetrics['image']etc.Additional featuresPlottingVisualization of metrics can be important to help understand what is going on with your machine learning algorithms.
Torchmetrics have built-in plotting support (install dependencies withpip install torchmetrics[visual]) for nearly
all modular metrics through the.plotmethod. Simply call the method to get a simple visualization of any metric!importtorchfromtorchmetrics.classificationimportMulticlassAccuracy,MulticlassConfusionMatrixnum_classes=3# this will generate two distributions that comes more similar as iterations increasew=torch.randn(num_classes)target=lambdait:torch.multinomial((it*w).softmax(dim=-1),100,replacement=True)preds=lambdait:torch.multinomial((it*w).softmax(dim=-1),100,replacement=True)acc=MulticlassAccuracy(num_classes=num_classes,average="micro")acc_per_class=MulticlassAccuracy(num_classes=num_classes,average=None)confmat=MulticlassConfusionMatrix(num_classes=num_classes)# plot single valueforiinrange(5):acc_per_class.update(preds(i),target(i))confmat.update(preds(i),target(i))fig1,ax1=acc_per_class.plot()fig2,ax2=confmat.plot()# plot multiple valuesvalues=[]foriinrange(10):values.append(acc(preds(i),target(i)))fig3,ax3=acc.plot(values)For examples of plotting different metrics try runningthis example file.Contribute!The lightning + TorchMetrics team is hard at work adding even more metrics.
But we're looking for incredible contributors like you to submit new metrics
and improve existing ones!Join ourSlackto get help with becoming a contributor!CommunityFor help or questions, join our huge community onSlack!CitationWe’re excited to continue the strong legacy of open source software and have been inspired
over the years by Caffe, Theano, Keras, PyTorch, torchbearer, ignite, sklearn and fast.ai.If you want to cite this framework feel free to use GitHub's built-in citation option to generate a bibtex or APA-Style citation based onthis file(but only if you loved it 😊).LicensePlease observe the Apache 2.0 license that is listed in this repository.
In addition, the Lightning framework is Patent Pending.