pytorch / ignite

High-level library to help with training and evaluating neural networks in PyTorch flexibly and transparently.
https://pytorch-ignite.ai
BSD 3-Clause "New" or "Revised" License
4.55k stars 619 forks source link

Customization of reduction/gathering ops for Metric #1288

Open vfdev-5 opened 4 years ago

vfdev-5 commented 4 years ago

🚀 Feature

Idea is to make configurable Metric's reduction/gathering ops. By default, we are using our code, but user can globally override those functions. For example, if uses a custom unsupported distributed framework, or deals with asymmetry like here etc

EDIT:

When a metric is implemented methods like reset, update and compute are decorated with reinit__is_reduced and sync_all_reduce. sync_all_reduce is implemented here: https://github.com/pytorch/ignite/blob/581f5b4841d68023de512de20757f42b55dc3105/ignite/metrics/metric.py#L550-L594 where we are using idist.all_reduce(t, **op_kwargs) So, the issue desctiption says:

Idea is to make configurable Metric's reduction/gathering ops. By default, we are using our code, but user can globally override those functions.

In other words, we would like to be able to call user custom all_reduce instead of idist.all_reduce


A tentative API for this feature


import ignite.distributed as idist
from ignite.metrics import set_all_reduce_fn, reset_all_reduce_fn, get_all_reduce_fn
from ignite.metrics import Accuracy

def my_all_reduce(tensor: Union[torch.Tensor, float], op: str = "SUM", **kwargs):
    # ... custom implementation
    pass

set_all_reduce_fn(my_all_reduce)
assert get_all_reduce_fn() == my_all_reduce

acc = Accuracy()
acc.update(...)
value = acc.compute()  # should call my_all_reduce

reset_all_reduce_fn()
assert get_all_reduce_fn() == idist.all_reduce
sdesrozis commented 4 years ago

A really nice feature