MONAI icon indicating copy to clipboard operation
MONAI copied to clipboard

Lazy eval refactor

Open atbenmurray opened this issue 1 year ago • 32 comments

Lazy evaluation refactor

This PR contains a reworking of the lazy evaluation mechanism on dev to provide the following:

  1. There are now three lazy execution options:
    1. the options=None (default) mode guarantees that all pending operations are evaluated before evaluating non-lazy _apply_transforms or transforms with lazy set to False. This means that all pipelines should be able to execute lazily without breaking (the first thing that users will try)
    2. there is an options={'reorder': 'lazy_last'} mode that moves all lazy operations past non-lazy ones, (bounded by use of ApplyPending/ApplyPendingd) and guarantees invertibility
    3. there is an options={'reorder': 'lazy_last_nosync'} mode that relies on ApplyPending/ApplyPendingd to cause application of pending transforms, and precisely mimics the behaviour of the lazy execution on dev.
  2. Lazy now has three states on Compose:
    1. lazy=True overrides the lazy flag on lazy transforms, making them all execute lazily
    2. lazy=None uses the lazy flags set on transforms, and can execute them lazily or non-lazily depending
    3. lazy=False (default) causes lazy processing to be disabled and all lazy transforms will execute non-lazily
  3. Lazy transforms have a lazy override that they can optionally set during __call__. This means that lazy behaviour can be overridden without mutating the transform instance
  4. Logging has been fixed so that the first transform is included in the logging messages. Logging activities are clearer and contain addition important information about lazy modes
  5. ApplyPending and ApplyPendingd replace the use of Identity and Identityd for causing pending ops to be evaluated when using the lazy_last[_nosync] options
  6. Compose API changes:
    1. overrides and override_keys have been consolidated into overrides
      1. on dev: override_keys=('a', 'b'), overrides={'mode': ('bilinear', None), 'padding_mode: (None, 'zeros')'}
      2. on PR: overrides={'a': {'mode': 'bilinear'}, 'b': {'padding_mode': 'zeros'}}
    2. log_stats and verbose have been replaced with logger_name, consistent with other parts of MONAI
  7. additional lazy resampling documentation has / is being added
  8. additional lazy tests have been added

Types of changes

  • Breaking changes:
    • lazy=True/None and options=None is new functionality relative to dev
    • lazy=True/None and options={"reorder": "lazy_last"} is new functionality relative to dev
    • lazy=None and setting lazy=True on a transform is new functionality relative to dev
    • Identity[d] doesn't cause pending operations to be evaluated; they must be replaced with ApplyPending[d]
  • Non-breaking changes:
    • lazy=True and options={"reorder": "lazy_last_nosync"} (with Identity[d]being replaced byApplyPending[d]) should be identical in behaviour to dev`
  • [wip] 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 --disttests.
  • [x] In-line docstrings updated.
  • [wip] Documentation updated, tested make html command in the docs/ folder.

atbenmurray avatar Mar 30 '23 09:03 atbenmurray

hi @atbenmurray sorry, earlier I didn't notice this PR is pending review. please address the merging conflicts I'll test and merge it. ApplyPending is better than implicit resampling calls in compose, and allows for flexible overriding.

wyli avatar Apr 06 '23 10:04 wyli

Gone through the merge from dev. I still need to test things before we merge!

atbenmurray avatar Apr 14 '23 16:04 atbenmurray

@wyli, @Nic-Ma, @ericspod I think that all of the tests should now be passing on this. I'm updating docs and adding a few additional unit-tests but I would say that, with a final review, this should be merge-ready once I've resolved conflicts

atbenmurray avatar Apr 24 '23 09:04 atbenmurray

@wyli, @ericspod I think this is about as good as it gets on passing tests, yes? Please can we try to get this in tomorrow? I'm not sure what this GPU fail is about

atbenmurray avatar Apr 25 '23 17:04 atbenmurray

  • needs more design discussions for monai/transforms/transform.py Lazy eval refactor #6257 (comment)
  • although the pre-merge quick tests work fine, all the relevant integration tests are failing, because of the API changes introduced.

I think the design in this PR represents a significant improvement over dev.

The failing tests are the ones using bundles yes? Can you tell me what I do to update them? I see the failure on the integration tests but I haven't used bundle before.

atbenmurray avatar Apr 27 '23 12:04 atbenmurray

https://github.com/Project-MONAI/MONAI/issues/6450 discusses compiler options. I think it this point it is our best bet.

atbenmurray avatar Apr 28 '23 18:04 atbenmurray

@wyli @ericspod Ok, so I've implemented the changes discussed as a result of Friday's meeting and subsequent conversations

Lazy resampling now has three modes, controlled by an options parameter on compose:

  1. "in_order": this preserves the ordering while executing during lazy mode: options = None and means that users pipelines cannot be broken by setting lazy=True
  2. "reorder: lazy_last": this allows some reordering but treats ApplyPending/ApplyPendingd as a barrier, so reordering can only happen across ApplyPending/ApplyPendingd. It supports inverse and works the same way as OneOf, RandomOrder, etc. To use it you supply options={'reorder': 'lazy_last'}
  3. "reorder: lazy_last_nosync": this should mimic the dev mode behaviour precisely. It allows effective execution order to be different for different keys in a dictionary and does not allow inverse. To use it, you supply options={'reorder': 'lazy_last_nosync'}

I've also tidied up some of the validation so that inverse behaviour is clearer to the user

Tests need augmenting still as does documentation

atbenmurray avatar May 02 '23 09:05 atbenmurray

  1. "reorder: lazy_last_nosync": this should mimic the dev mode behaviour precisely. It allows effective execution order to be different for different keys in a dictionary and does not allow inverse.

the current dev branch supports invert, which is used in benchmarking the CT segmentation workflow https://github.com/Project-MONAI/model-zoo/pull/397/ please see the train.json file in that PR. without the inverse, we may not be able to achieve the same validation scores. please double-check.

wyli avatar May 02 '23 10:05 wyli

When does it support it. In the case where there are no non lazy transforms?

atbenmurray avatar May 02 '23 11:05 atbenmurray

@wyli Please confirm for me the following:

As I understand it, you can invert with lazy_last_nosync if and only if your pipeline contains no transforms that are both invertible and not lazy (ie Lambda)

atbenmurray avatar May 02 '23 11:05 atbenmurray

@wyli Please confirm for me the following:

As I understand it, you can invert with lazy_last_nosync if and only if your pipeline contains no transforms that are both invertible and not lazy (ie Lambda)

yeah, I'm not sure about the "if and only if" statement. "pipeline contains no transforms that are both invertible and not lazy" is sufficient but not necessary.

wyli avatar May 02 '23 11:05 wyli

@wyli Please confirm for me the following: As I understand it, you can invert with lazy_last_nosync if and only if your pipeline contains no transforms that are both invertible and not lazy (ie Lambda)

yeah, I'm not sure about the "if and only if" statement. "pipeline contains no transforms that are both invertible and not lazy" is sufficient but not necessary.

I'll rephrase: . If you have transforms that are not invertible, they aren't added to the applied operations, so are not touched during inverse . If you have only lazy transforms and ApplyPendings, again you are fine because the lazy transforms are ordered correctly . If you have transforms that are invertible but not lazy, they mean that the applied order is not compatible with the transform order, and so inverse cannot be run

Hence if and only if contains no non-lazy, invertible transforms

atbenmurray avatar May 02 '23 12:05 atbenmurray

@wyli Please confirm for me the following: As I understand it, you can invert with lazy_last_nosync if and only if your pipeline contains no transforms that are both invertible and not lazy (ie Lambda)

yeah, I'm not sure about the "if and only if" statement. "pipeline contains no transforms that are both invertible and not lazy" is sufficient but not necessary.

I'll rephrase: . If you have transforms that are not invertible, they aren't added to the applied operations, so are not touched during inverse . If you have only lazy transforms and ApplyPendings, again you are fine because the lazy transforms are ordered correctly . If you have transforms that are invertible but not lazy, they mean that the applied order is not compatible with the transform order, and so inverse cannot be run

Hence if and only if contains no non-lazy, invertible transforms

Still inaccurate... if there's Identity[D] transforms placed accordingly before the non-lazy invertible transforms, the current dev branch can do the inverse.

If you are looking for the precise condition of not allowing inverse, it's been described in this PR in the warning message https://github.com/Project-MONAI/MONAI/pull/6372/files

inversing of a sequence of user-defined transforms applied to a metatensor is not possible if, when applying the transforms, applied_operations of that metatensor is modified while its pending_operations is not empty.

wyli avatar May 02 '23 12:05 wyli

@wyli Please confirm for me the following: As I understand it, you can invert with lazy_last_nosync if and only if your pipeline contains no transforms that are both invertible and not lazy (ie Lambda)

yeah, I'm not sure about the "if and only if" statement. "pipeline contains no transforms that are both invertible and not lazy" is sufficient but not necessary.

I'll rephrase: . If you have transforms that are not invertible, they aren't added to the applied operations, so are not touched during inverse . If you have only lazy transforms and ApplyPendings, again you are fine because the lazy transforms are ordered correctly . If you have transforms that are invertible but not lazy, they mean that the applied order is not compatible with the transform order, and so inverse cannot be run Hence if and only if contains no non-lazy, invertible transforms

Still inaccurate... if there's Identity[D] transforms placed accordingly before the non-lazy invertible transforms, the current dev branch can do the inverse.

If you are looking for the precise condition of not allowing inverse, it's been described in this PR in the warning message https://github.com/Project-MONAI/MONAI/pull/6372/files

inversing of a sequence of user-defined transforms applied to a metatensor is not possible if, when applying the transforms, applied_operations of that metatensor is modified while its pending_operations is not empty.

Identity/IdentityD will need to be replaced with ApplyPending/ApplyPendingd, but fine, yes, they are also not added to the applied operations so don't break invertibility

The problem with the implementation as it is is that the warning describes the symptom rather than the cause when running forward. When running inverted, the error that gets raised is no more helpful:

RuntimeError: Error Lambda getting the most recently applied invertible transform Rotate 140236029348832 != 140236029348880

The error should explain that the pipeline that has been created cannot be executed in inverse and ideally should give a clear indication as to why. I'm trying to find the best way to do that, but it should really be clear from a static analysis of the transform list rather than not catching the condition and waiting for an internal mechanism to raise an exception

atbenmurray avatar May 02 '23 12:05 atbenmurray

@wyli I think the best approach is to add a flag to the metadata on the tensor if we are running out of order and there are pending operations when an operation is added to applied. That way, we can simply check the tensor before we invert and provide a meaningful error message

Edit: Looks like the warning is already added to the info so this should just be a tweak to the existing code.

atbenmurray avatar May 02 '23 17:05 atbenmurray

currently the runtime error will be raised with an additional warning of the details when accessing the applied_operations in such cases: https://github.com/Project-MONAI/MONAI/blob/4bfc8e97661d44cc13553721e5025005757bad91/tests/test_invert.py#L100-L102

wyli avatar May 02 '23 19:05 wyli

I think it is better to have inverse throw the reasons why data isn't convertible. Those warnings tend to get lost in the noise. Given that we capture the reason why inverse will fail during the forward pass, we really ought to make use of it if we know inverse will fail. It's much clearer than the Id mismatch error

Anyway, that should avoid the use of a can invert flag on the execution policy, so it's definitely an improvement from that flag. Isnt it now equivalent?

atbenmurray avatar May 02 '23 20:05 atbenmurray

test_integration_lazy_samples.py now passes for option=None and option={"reorder": "lazy_last_nosync"}

atbenmurray avatar May 03 '23 12:05 atbenmurray

test_integration_lazy_samples.py now passes for option=None and option={"reorder": "lazy_last_nosync"}

Hi @atbenmurray there are min-tests failing... please let us know when this is ready for review. (We'll also need a few days for providing comments and addressing comment about the code changes and repeating some basic benchmarks to confirm.)

wyli avatar May 03 '23 14:05 wyli

test_integration_lazy_samples.py now passes for option=None and option={"reorder": "lazy_last_nosync"}

Hi @atbenmurray there are min-tests failing... please let us know when this is ready for review. (We'll also need a few days for providing comments and addressing comment about the code changes and repeating some basic benchmarks to confirm.)

From a code standpoint, this is ready for review. I wouldn't worry about the documentation; if code changes are made it might also change the docstrings for the classes involved. My intention is to create a new document in the docs folder that deals with the concepts and design of lazy resampling in one place, and link to it from all of the classes involved, so there is less repetition and the docs for the individual classes can be more succinct about it

atbenmurray avatar May 04 '23 08:05 atbenmurray

we need basic docstrings about how to use it, and please update the PR description so that we are on the same page and we can start running the benchmarks.

Goal of this PR: Screenshot 2023-05-04 at 09 56 10 https://github.com/Project-MONAI/MONAI/issues/6450#issuecomment-1529003032

@ericspod @Nic-Ma @KumoLiu please also help review

wyli avatar May 04 '23 08:05 wyli

PR text is updated to list the changes

atbenmurray avatar May 04 '23 12:05 atbenmurray

PR text is updated to list the changes

can you highlight for each item whether they are breaking changes or non-breaking new features?

wyli avatar May 04 '23 13:05 wyli

Having an issue with tests/test_integration_lazy_samples.py It does inversion by pulling out the transforms from the Compose object (this is something that we should really not be doing). You can never do it with OneOf, SomeOf, or RandomOrder Those transforms are reordered by options={"reorder": "lazy_last"}, obviously the test then fails. I think the test should be updated to use the compose.inverse instead. Is there a specific reason that it isn't?

atbenmurray avatar May 04 '23 13:05 atbenmurray

I had another round of reviewing this PR (up to commit b151607), now it introduces multiple new concepts such as compose compiler and lazy executors, I think the code design requires more discussions, in my opinion, we should think about code flexibility, readability, future extensions, I'm not sure if the current implementation is better than the current dev branch, it may take substantial resources to discuss, revise, verify.

'dev' branch default behaviour breaks general case pipelines. As a first release of lazy resampling, I'd argue that it makes it not fit for purpose and most people shouldn't try it.

atbenmurray avatar May 04 '23 14:05 atbenmurray

ComposeCompiler is a design that I've tested out on the full branch. You've seen it before. In this case, it is a very small class that, given a flag, produces a policy object. It is also docstringed as not being part of the public API.

The lazy execution code has been separated out into 'in_order' and 'out_of_order' execution. They are no less tested (actually more tested) than what is on dev.

Reordering of the transforms uses the same mechanism as is used by OneOf, SomeOf and RandomOrder.

All of this was necessary in order to provide an optional lazy mode that matched dev for the benchmarks and a default mode that doesn't break pipelines in the general case.

atbenmurray avatar May 04 '23 14:05 atbenmurray

@wyli At this point, I think that the best thing to do is to try to propose the most constructive way forward that I can think of.

I think that the PR is very close to being ready. @ericspod has gone over the PR several times and also believes that the solution is close to merge ready.

The primary outstanding task from a correctness standpoint is to confirm that lazy=True, options={"reorder": "lazy_last_nosync"} does what dev does. That is something that you are best placed to do.

The remaining tasks are making sure we haven't missed anything from a test standpoint and documenting.

If you have design concerns, I am more than happy to meet with you and discuss them properly. We know that there are documentation holes, but documentation is low risk and, in any case, is something that I have been working on today. Similarly, we can ensure that we haven't left any testing holes, either left over from dev or newly introduced in the PR.

ComposeCompiler in particular is an implementation detail. I've discussed the design with you and presented it at several dev meetings, but it is documented as being an implementation detail that users are not supposed to interact with directly, so it can be subsequently changed without impact.

the different lazy executors have been added by necessity. The "in order" executor implements what was supposed to be the default behavour that has been discussed and presented throughout. The "out of order" executor carefully mimics the behavour on dev and I believe it now does that completely. If you can find any issues with it, we can certainly deal with them. I separated them out from each other because I found the cyclomatic complexity of evaluate_with_overrides to be prohibitive even as it was, and trying to have those two different execution strategies bound up in the same code would have caused many problems, especially if we then required a third. This design separates them out and makes the logic of when one needs to apply pending transforms and when one doesn't to be very clear.

We can discuss any of these matters directly if you prefer. I am happy to meet, as always.

atbenmurray avatar May 04 '23 16:05 atbenmurray

sure @atbenmurray, our discussions haven't been very productive/constructive... and I don't think I have enough energy to follow up or to write more comments about this PR. (but as one of the core maintainers of this codebase, I wholeheartedly hope this PR is not merged in the current form.)

wyli avatar May 04 '23 17:05 wyli

Hi @atbenmurray @ericspod @wyli ,

Sorry I didn't follow carefully in this PR, currently, I think the biggest gap of design discussion is the transform order? And this PR introduced many new concepts to solve the order problem in new ways, to help quickly move forward, how about let's do a simplified step first: We only do lazy resampling for continuous spatial / crop pad transforms, then we can guarantee the order of other transforms? I think we don't need much code change to achieve it in MONAI 1.2. When this first step merged, we can enhance it with more fancy methods to adjust the transform order in the next version? What do you think? Please feel free to correct me if I misunderstood anything.

Thanks in advance.

Nic-Ma avatar May 04 '23 17:05 Nic-Ma

Hi @atbenmurray @ericspod @wyli ,

Sorry I didn't follow carefully in this PR, currently, I think the biggest gap of design discussion is the transform order? And this PR introduced many new concepts to solve the order problem in new ways, to help quickly move forward, how about let's do a simplified step first: We only do lazy resampling for continuous spatial / crop pad transforms, then we can guarantee the order of other transforms? I think we don't need much code change to achieve it in MONAI 1.2. When this first step merged, we can enhance it with more fancy methods to adjust the transform order in the next version? What do you think? Please feel free to correct me if I misunderstood anything.

Thanks in advance.

Hi Nic, the problem that we would have is that we need to replicate the dev behaviour in order to avoid the new QA cycle. Unless we have the two modes (options=None and options={'resize': 'lazy_last_nosync'}) that have been implemented in this PR we are back to the same issues:

  1. Without the default, no reordering mode, the pipelines will break / output wrong results in the general case: [Flip, Spacing, Rand3DElastic, Rotate] becomes [Rand3DElastic, Flip, Spacing, Rotate]. This will upset a lot of users, I think
  2. Not having 'benchmark' mode {'rotate': 'lazy_last_nosync'} will invalidate the benchmarks and impact their performance, but it is absolutely only for expert users as they have to synchronise with ApplyPending[d] correctly or they will get exceptions / a different pipeline

I've had a week with the dev behaviour and I believe that I have replicated it successfully, barring issues that arise during the integration tests. I've also provided the standard lazy resampling mode as per the design that does no reordering and should be completely safe with all pipelines. Nothing should break when the user sets lazy=True.

We agreed last Friday that this was probably the best way forward and the discussion is reflected in #6450. Providing the safe (no-out of order execution) mode as per the design means users can set lazy=True safely. Providing the dev behaviour means no benchmark invalidation.

I went over the design and implementation with @ericspod today; I'm happy to go over the design and implementation with you and/or @wenqi tomorrow.

I'd honestly be more worried about changing what is on dev presently. evaluate_with_overrides is a highly complex function, firstly. Secondly, it's so easy to overlook edge cases (see dev and reordering of valid pipelines) that we'd have to live with that new mode for a week and think about what supporting it means moving forward. At least the modes that we have are conceptually well understood.

I've done everything that I can to minimise the risks and I would definitely worry about the risks of yet another lazy execution definition change.

atbenmurray avatar May 04 '23 18:05 atbenmurray