[WIP] 2220 Decollate batch into list of tensors after model forward

[Nic-Ma]

Description

This PR added support to handle list of Tensor in post transforms.

Status

Work in progress

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[Nic-Ma]

Hi @ericspod @wyli @rijobro ,

In this drat PR, I tried to enhance the Activationd transform to support a list of Tensor.
If you have no concern about this solution, I will try to enhance all the other post transforms and utility transforms.
Or if you have any better idea, please feel free to share some sample code here.

Thanks in advance.

[ericspod]

It looks fine to me. Do we have a particular use case for this?

[Nic-Ma]

It looks fine to me. Do we have a particular use case for this?

Hi @ericspod ,

Thanks for your review.
When we try to invert the transforms, the results usually get a list of channel-first Tensor with different size. So we need to support it in the following transforms. For example:

post_transforms = Compose([
    Activations(keys="pred"),
    CopyItemsd(keys="pred", times=1, names="inverted_pred"),
    Inverted(keys="inverted_pred", nearest_interp=False),
    AsDiscreted(keys=["pred", "inverted_pred"]),
    SaveImaged(keys=["pred", "inverted_pred"]),
])

Thanks.

[ericspod]

@rijobro mentioned abstracting the concept out rather than inserting it into every transform everywhere. Earlier we had discussed transforms representing diverging paths in the transform sequence, for example something like OneOf where one of the transforms passed to the object as an argument is randomly chosen to be applied. Perhaps a wrapping transform that takes a list of tensors as you want and applies each in the list to the sequence of transforms it was given then either returns the list or collates the list. This would look like:

post_transforms = Compose([
    ApplyToSeqd(keys="pred", collate=False, [
        Activations(keys="pred"),
        CopyItemsd(keys="pred", times=1, names="inverted_pred"),
        Inverted(keys="inverted_pred", nearest_interp=False),
        AsDiscreted(keys=["pred", "inverted_pred"]),
        SaveImaged(keys=["pred", "inverted_pred"]),
    ])  # output is dictionary with lists of tensors for values
])

[Nic-Ma]

Hi @rijobro and @ericspod ,

So in summary, now we have 4 draft ideas to solve this problem:

1. Add recursive call in every post transform to handle a list of tensor (Nic)
2. Extract (1) to the base class, so don't need to change every transform, haven't got how to implement it (Richard)
3. Decollate a batch tensor to a list of tensor, then all the transforms and metrics handle every tensor in a for-loop (Wenqi)
4. Make a wrapper transform to handle the list of tensor, also can support other cases like OneOf (Eric)

Thanks.

[rijobro]

Good summary, let's discuss tomorrow.

[Nic-Ma]

Thanks for this great online discussion.
We all agreed to go with decollate way.
It's a big change to many modules, I will try to implement it soon.

Thanks.

[wyli]

Thanks for this great online discussion.
We all agreed to go with decollate way.
It's a big change to many modules, I will try to implement it soon.

Thanks.

thanks, if needed please create a branch in the codebase dedicated to this feature, it'll make it easier to manage the release of this feature

[Nic-Ma]

Hi @wyli ,

Thanks for your suggestions.
I already split the task into steps.
If can't clearly submit PR for steps, I will create a branch.

Thanks.

[Nic-Ma]

I already updated all the necessary post / IO / utility transforms to be channel-first, will try to update MONAI engines and handlers in the next steps.

Thanks.

[wyli]

I already updated all the necessary post / IO / utility transforms to be channel-first, will try to update MONAI engines and handlers in the next steps.

Thanks.

thanks, I think we'll need full integration tests before merging this feature... these pre-merge tests won't be reliable enough for this PR.

[Nic-Ma]

Hi @wyli ,

I will trigger integration tests when all the things are ready.

Thanks

[Nic-Ma]

Hi @wyli @ericspod @rijobro ,

I am adjusting the metrics interface, after decollating, there are several options:

1. Input a list of channel-first Tensor for y_pred and y.
2. Input only 1 channel-first Tesnor for y_pred and y, do for-loop for 1 batch outside metrics, just same as transforms. But some metrics may need to compute on a batch data or the whole data together, like AUC, etc.

I prefer to unify the metrics API to below signature, which was discussed in PR #565:

class Metric:
    def update(self, y_pred: Union[torch.Tensor, Sequence[torch.Tensor]], y: Union[torch.Tensor, Sequence[torch.Tensor]]):
        # call for a decollated batch data usually when iteration completed
        # compatible with `batch-first tensor` and `list of channnel-first tensor`
        for pred_, y_ in zip(y_pred, y):
            self._add_sample(pred_, y_)
    
    def _add_sample(self, y_pred: torch.Tensor, y: torch.Tensor):
        # execute logic for 1 sample of the batch
        pass

    def compute(self):
        # call for final computation usually when epoch completed
        pass

What do you guys think?

Thanks.

[Nic-Ma]

After a quick online discussion, I would prepare a separate PR to unify the basic metrics API.
It should support both batch-first Tensor and list of channel-first Tensor.
And will close ticket #561 #497 .

Thanks.

[Nic-Ma]

As this is a big feature update, I created a branch feature/2220-decollate in Project-MONAI repo.
I will merge this PR into that branch and create a new PR to dev branch.

Thanks.