meta tensor

[rijobro]

Description

First PR to MetaTensor branch.

As promised, we'll merge to MetaTensor branch via PRs and code reviews. This is the first PR and just adds the object and some tests. MetaObj allows for subclassing np.ndarray in the future should we want to.

Status

Ready

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• New tests added to cover the changes.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[wyli]

This is non-breaking, I think we should try to merge this into the dev branch directly so that we can potentially get more early feedback... after this PR we can collaborate on the MetaTensor branch at full speed for all the necessary updates.

for the naming of key classes and modules: meta_obj, MetaObj, meta_tensor, MetaTensor, do you agree with the current PR? @ericspod @Nic-Ma @atbenmurray @yiheng-wang-nv (can discuss on Friday as well).

[rijobro]

Should MetaTensor be MetaTensorImage to add in another layer of inheritance and allow for MetaTensorPoints etc?

[rijobro]

@wyli Is there much point merging this into dev? It won't be used until we have some functionality that uses it.

[wyli]

@wyli Is there much point merging this into dev? It won't be used until we have some functionality that uses it.

yes it'll be convenient for creating demos or experimenting with metatensor using pip install monai-weekly or pip install from git+dev branch. I think the MetaTensor branch will be largely broken as we introduce breaking changes in the following weeks/months.

[Nic-Ma]

Hi @rijobro ,

Thanks for your initial cool idea and quick PR.
If plan to merge to dev branch, I have 3 minor concerns:
(1) As this is a very important base classes, I feel we should write very detailed doc-string with example code and explanation for every arg. Refer to transform:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/transforms/transform.py
(2) MetaObj is the base class for Metadata, it should not be aware of any tensor or numpy specific properties.
(2) Let's try to define fewer APIs in the MetaObj base class to keep it simple and easy to extend.

What do you think?

Thanks.

[Nic-Ma]

Hi @vfdev-5 ,

We are trying to extend PyTorch Tensor with medical metadata properties.
Could you please help also take a look at this first PR from the PyTorch team angle?

Thanks in advance.

[vfdev-5]

@Nic-Ma thanks for pinging !

I'll check the PR in details a bit later, right now few comments/info (in case if helpful):

• In torchvision there is a work on inheriting from pytorch tensors => tensor images:
  https://github.com/pytorch/vision/blob/08cc9a7fb5baa9655fa5d2dd44a6fcfcbe48a1c7/torchvision/prototype/features/_image.py#L37 , you can check out how it is done
• "meta" name may be confusing as pytorch has "meta" device as well (~ lazy tensors)

[Nic-Ma]

And as @vfdev-5 mentioned in another conversation, this MetaTensor may have issue with TorchScript?

Thanks.

[rijobro]

Hi, think I've tried to resolve any outstanding issues. I also added testing for pickling, torchscript and amp.

[rijobro]

@wyli @Nic-Ma I can't reproduce this error anywhere and I've tried on quite a few different machines (MacOS, Windows, 3xUbuntu). Can any of you reproduce?

[wyli]

@wyli @Nic-Ma I can't reproduce this error anywhere and I've tried on quite a few different machines (MacOS, Windows, 3xUbuntu). Can any of you reproduce?

I can replicate with MacOS py3.8 and with pip install -U -r requirements-dev.txt, let me have a look...

[rijobro]

@wyli let me know when you've bumped the min pytorch version, hopefully this PR will be good to go at that point.

[rijobro]

Hi @vfdev-5 have you had the chance to look this PR over? Would be great to get your insight.

[Nic-Ma]

Hi @vfdev-5 have you had the chance to look this PR over? Would be great to get your insight.

Hi @rijobro ,

I synced with @vfdev-5 offline, he tried to review this PR several times, but is still not clear about the context and initial intent of this PR, could you please help provide some design doc or usage examples to help him better understand it?

Thanks in advance.

[rijobro]

Sure @Nic-Ma. Hi @vfdev-5, we would like to attach our meta data to our images. Looking at the pytorch docs for extending torch.Tensor, ours is similar in concept to the MetadataTensor, but in implementation it more similar to the LoggingTensor example. This is because we subclass torch.Tensor rather than wrap it so that isinstance still works.

Previously, we used dictionaries of data containing images, their metadata and the transforms that have been applied to them, e.g., {"img": torch.Tensor, "img_meta_dict": dict, "img_transforms": list}. Throughout our code, we sometimes hardcode the suffix _meta_dict, and in other places we allow the user to change the default. This gets messy, especially when we have multiple images (e.g., img and seg).

Looking to the future, we would like our transforms to be able to update metadata when they are applied, so we need to first tidy things up, as there isn't scope for adding extra functionality as things stand.

Hence, what would have previously been:

transform = LoadImaged("img", "seg")
data = {"img": some_filename, "seg": some_other_filename}
data = transform(data)
print(data)  #  {"img": ...,  "img_meta_dict": ..., "seg": ..., "seg_meta_dict": ...}

is now

print(data)  # {"img": ..., "seg": ...}

and now "img" has the meta data and affine transformation matrix that can be extracted with:

im = data["img"]
print(im.meta)  # dictionary with meta data
print(im.affine)  # 4x4 or 3x3  torch.Tensor

Now that every image has its attached affine transformation matrix, it should be much easier to make transforms that correctly update it.

We allow users to disable this functionality in case they aren't interested in it, and it is causing them bother for whatever reason.

[vfdev-5]

@rijobro thanks a lot for details!

In terms of implementation, I'll check and comment out in the code.

Another point I was wondering as far as I understand tensor subclassing feature but not sure if this point is related to your idea. When subclassing torch tensor with __torch_function__ we have a possibility to route depending on input func whether ops are changing output type (back to torch.Tensor) or keep it as a subclass.
So, in your case looks like any op keep MetaTensor as MetaTensor. For example, if we sum up it we are still in the domain:

t = torch.tensor([1,2,3])
affine = torch.eye(4)
meta = {"some": "info"}
m = MetaTensor(t, affine=affine, meta=meta)
res = m.sum()
print(res)
> tensor(6)
MetaData
	some: info
	affine: tensor([[1., 0., 0., 0.],
        [0., 1., 0., 0.],
        [0., 0., 1., 0.],
        [0., 0., 0., 1.]])

If this is expected and intended behaviour, please ignore this point.

EDIT:
I was also wondering about how would you merge meta data for ops between several MetaTensors. Right now left operand is defining the meta data and this could a bit counter intuitive for commutative ops like add, multiply etc:

t = torch.tensor([1, 2, 3])
affine = torch.eye(4)
meta = {"some": "info"}

m1 = MetaTensor(t, affine=affine, meta=meta)

t = torch.tensor([10, 20, 30])
affine = torch.eye(5)
meta = {"some1": "info2"}

m2 = MetaTensor(t, affine=affine, meta=meta)

r1 = m1 + m2
r2 = m2 + m1
# r1.meta != r2.meta

[wyli]

When subclassing torch tensor with __torch_function__ we have a possibility to route depending on input func whether ops are changing output type (back to torch.Tensor) or keep it as a subclass.

that's a good point, perhaps we need this to differentiate w/o the batch dim for the collate... （the collate related feature is not included in this PR, we can consider this in the follow-ups.)

[wyli]

/build

[rijobro]

@vfdev-5 thanks for the useful input.

The way that we currently have it, yes, if torch.Tensor__torch_function__ would return a torch.Tensor (which is most operations, barring things like __repr__), then our default implementation is to return a MetaTensor.

So MetaTensor([1]) + torch.tensor([1]) will return the type MetaTensor. This is by design, but can be disabled with set_track_meta(False).

Also, yes, things get a bit confusing when combining meta data. But if m1 is CT and m2 is MRI, then what does it really mean to do m1 + m2 – what should the meta data be? Thankfully, I don't think we'll be doing that sort of operation often. If we do it at all, more likely we'll be doing some sort of operation between image and segmentation, in which they should (a) have similar meta data, and (b) the image would normally come first and so should have its meta data propagated.

[rijobro]

@wyli for collating and decollating, I've been messing around with inspect. If the function is torch.stack and the parent function contains collate, then do collation stuff. If the function is torch.unbind and the parent function contains decollate do decollation stuff.

This would be a follow up PR.

What do you think? It's maybe a little hacky. The alternative is to put that logic into the collation and decollation functions, which is what I did in the original PR.

[wyli]

thanks, it looks good to me, I put some minor docstring updates here rijobro#3

[wyli]

/build