Allow general pickle-able types as state objects?

[christopherbate]

🚀 Feature + Motivation

Some types of metrics, such as "COCO detection metrics" for computer vision are more easily computed by operating on lists of dictionaries containing arrays rather than pure tensors. In this case, the developer would like to keep multiple pieces of information about each prediction, such as the example ID, bounding boxes and associated scores, etc. I recently wrote a torchmetrics wrapper for COCO API metrics, but it was made more complicated by requiring the state objects to be tensors. It would be trivial if I could keep a list of dictionaries as the state object. Synchronization in Torch >= 1.8.0 can be done easily via all_gather_objects but for < 1.8.0 requires some pickling into a Tensor hack.

Pitch

Allow state objects to have lists of pickle-able objects.

Alternatives

Without this, I need to keep several state objects broken out by the sub-keys of my dictionary. It's not necessarily difficult but it is more code.

I can probably help do this unless someone can point out why it's too difficult/impossible. I still need to dig into the source a bit more.

[github-actions]

Hi! thanks for your contribution!, great first issue!

[zhiqwang]

Hi @christopherbate ,

torchvision have also added a all_gather on arbitrary picklable data here that supporting lower version PyTorch.

BTW I've also written a quick torchmetrics wrapper for COCO API metrics COCOEvaluator based on torchvision's CocoEvaluator. But I didn't test the distributed mode, and additional checks and further encapsulation are needed here.

[SkafteNicki]

Hi @christopherbate,
I would be fine with this improvement :]
To make it easy we can enforce that anything other than empty list and tensors are only allowed if you have pytorch v1.8 or higher.

[christopherbate]

Hi @christopherbate ,

torchvision have also added a all_gather on arbitrary picklable data here that supporting lower version PyTorch.

BTW I've also written a quick torchmetrics wrapper for COCO API metrics here based on torchvision's CocoEvaluator. But I didn't test the distributed mode, and additional checks and further encapsulation are needed here.

That's interesting, thanks.
I used detectron2's implementation as a reference since I'm more familiar with their API and trust its distributed implementation.

Regarding the all_gather operation, then, it should be easy enough to borrow that one method from torchvision, thanks for pointing that out.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.