Adding components and refactoring of schedulers

[ericspod]

Addresses #280.

The objective here is define a base class for all schedulers that we can use for reducing duplication of code and type checking.

One thing added is a way of adding new noise schedule functions by adding them to a component store object. This version adds the cosine schedule but others can be added by users in their own scripts.

[marksgraham]

Hi @ericspod

I'm a fan of the idea in general, it makes sense to factor a lot of the noise-related parts out of each individual scheduler. Questions/thoughts:

1. I think it would be really great if people are able to play with their own custom noise schedules without having to commit them. I've recently been playing with funky noise schedules, but nothing I think would be useful enough to others that I would want to push to the codebase. As it stands I don't think this is possible, right? If I try to add my own custom schedule I need to modify scheduler.py?

2. A related general thought - at the moment we force noise schedules to be designed by setting beta, and we calculate alpha, alpha_cumprod from beta. It is sometimes useful to define the noise schedule in terms of the other params. For example, if you want a noise schedule that alters the SNR as a function of t, it is easiest to define in terms of alpha_cumprod which is a direct function of the SNR (discussed here). Do you think it would be useful to have some helper functions so users can write noise schedules in terms of alpha or alpha cumprod and we can return the corresponding beta values? (Alpha to beta is easy, going from alpha_cumprod to beta isn't hard but a bit fiddly as you have to reverse the cumulative product).

[marksgraham]

typo, should be "epsilon"

[ericspod]

Hi @ericspod

Thanks for looking at this @marksgraham!

The idea with the BetaSchedules object is actually to allow users to add their own. They can create their own functions and decorate them with @BetaSchedules.add_def which adds the functions to the collection. This is something like what we use in MONAI to do layer factories. There should be an example of what this would look like but it would like what's in the file now, so long as the decorated function is defined before it's needed this will work. This comes back to the questions I'm having about the design, is this too complicated versus passing a callable to the Scheduler class or doing something else? Is it flexible enough since you mentioned defining schedules in terms of something other than beta. The defined cosine and sigmoid functions ignore the beta start and end arguments but still work, do we want something that's flexible enough to allow for other arguments?

What it comes down to is that we need an idea of what is common to all noise schedulers and define our base class with that then specialise others, so consider what to define to make the designs modular and flexible to account for the variations you're describing.

[marksgraham]

Hi @ericspod

Ah I see- I had tried to add my own scheduler and failed, but it was because I was missing the decorator.

Regarding flexibility - I think this should be flexible enough in terms of return type. As I said users will either want to define in terms of alpha, beta, or alpha cumprod, but I think it will be sufficient to provide helper functions to go to alpha -> beta or alpha_cumprod -> beta, so the user can define the noise in whatever set of parameters they want but return beta. We could alternatively require that the noise schedule return alpha, beta, and alpha_cumprod, rather than having the scheduler calculate alpha/alpha cumprod from beta. Regarding arguments - It would be good if we were able to more flexibly allow for other arguments, without forcing the user to leave in an unused beta start/beta end which feels a bit clunky.

In terms of whether this is over-engineered: it certainly took me some time to get my head around the code. But I think it will be pretty easy for users to use in practice, if we supply examples. But one thing I'm unclear on is what advantage this offers over allow the user to pass a callable to the scheduler?

Tagging @Warvito in case he has any opinions on this.

[ericspod]

Hi @marksgraham,

I need to update this PR with this feedback in mind and we'll come around to it again, if beta/alpha/alpha_cumprod are all we need to worry about we can define functions to return these, or if only one value is returned assume this is beta and calculate the others from it as it currently is.

But one thing I'm unclear on is what advantage this offers over allow the user to pass a callable to the scheduler?

We could do that easily enough, my concern in MONAI and elsewhere is to make explicit what functions are available and give them useful names. We have enumerations for int and string constants in MONAI for this reason as opposed to the way other libraries take of just stating valid values in the docstring. Our layer factories make it easy to choose a layer by name with the dimensionality being parameter, but they also gives clear names to the concepts they store. Explicit is better than implicit, but simple is better than complex so the question is if the tradeoff makes sense here.

[marksgraham]

Hi @ericspod
Sorry for taking so long, now GTC+MICCAI are out the way I'll be more responsive.

We could do that easily enough, my concern in MONAI and elsewhere is to make explicit what functions are available and give them useful names. We have enumerations for int and string constants in MONAI for this reason as opposed to the way other libraries take of just stating valid values in the docstring. Our layer factories make it easy to choose a layer by name with the dimensionality being parameter, but they also gives clear names to the concepts they store. Explicit is better than implicit, but simple is better than complex so the question is if the tradeoff makes sense here.

OK, then I think the tradeoff does make sense here! I guess the next step is to update all the other schedulers? In which case I'll hold off carefully reviewing the latest batch of changes until thats done?

[ericspod]

Hi @marksgraham, I have things updated finally if you want to take a look.

[ericspod]

I have the conflict from the scheduler addition to fix however.

[marksgraham]

Hi @ericspod
I've taken a look and looks good - shall I hold off until you've fixed the schedule to do a formal review?
NB that the diffusion tutorials will have to be updated. Because you've kept the defaults the same it should be sufficient to just update the arguments, I don't think youll need to rerun all the notebooks.

[ericspod]

Thanks @marksgraham I think it's ready now, I looked through the tutorials and made changes where needed, though yes I haven't rerun them which raises the question of what to do for CI/CD with these.

[ericspod]

@Warvito Updates from comments added!

[marksgraham]

Hi @ericspod

This looks good. Just need to run ./runtests.sh --autofix as it picks up some formatting corrections

[Warvito]

I am getting these flake8 errors when running ./runtests.sh -f -u --net

flake8
6.0.0 (flake8-bugbear: 22.12.6, flake8-comprehensions: 3.10.1, flake8-executable: 2.1.2, mccabe: 0.7.0, pep8-naming: 0.13.3, pycodestyle: 2.10.0, pyflakes: 3.0.1) CPython 3.10.6 on Linux
/media/walter/Storage/Projects/GenerativeModels/generative/networks/schedulers/ddim.py:89:30: F541 f-string is missing placeholders
/media/walter/Storage/Projects/GenerativeModels/generative/networks/schedulers/pndm.py:92:30: F541 f-string is missing placeholders
/media/walter/Storage/Projects/GenerativeModels/generative/networks/schedulers/scheduler.py:35:1: F401 'typing.Callable' imported but unused
/media/walter/Storage/Projects/GenerativeModels/generative/networks/schedulers/scheduler.py:37:1: F401 'numpy as np' imported but unused
/media/walter/Storage/Projects/GenerativeModels/generative/utils/__init__.py:16:1: F403 'from .misc import *' used; unable to detect undefined names
2     F401 'typing.Callable' imported but unused
1     F403 'from .misc import *' used; unable to detect undefined names
2     F541 f-string is missing placeholders

[ericspod]

@Warvito I've made those fixes but my version of the tools brought up some other issues, different versions of flake8 keep coming up with new issues not seem before.

[marksgraham]

Do we need to settle on a version of flake8 to use then?

[ericspod]

We should always use the current version I think, it just means we'll encounter errors unrelated to our own changes on occasion.