[Diff] sbAsma/issue1279 noise conditioning #1358
Conversation
This commit resolves architectural incompatibilities when integrating
diffusion-based forecast engines:
1. FSDP Sharding: DiffusionForecastEngine wraps ForecastingEngine
as `self.net`, but trainer code assumed direct `fe_blocks` access. Fixed by:
- Adding fe_diffusion_model conditional check in init_model_and_shard()
- Routing to model.forecast_engine.net.fe_blocks for diffusion mode
2. Model Initialization: Reordered ForecastingEngine creation to handle both
standard and diffusion-wrapped variants with proper fallback.
3. Target Format Handling: Autoencoder mode uses different target
structure than diffusion mode. Added conditional formatting:
- Diffusion: targets = {"targets": [targets], "aux_outputs": aux}
- Autoencoder: targets = {"physical": batch[0]}
4. Config Updates: added file config/diffusion_config.yml for diffusion
model config
|
assign me |
sbAsma
changed the title
src/weathergen/model/model.py
Outdated
| ) | ||
| else: | ||
| self.forecast_engine = ForecastingEngine(cf, self.num_healpix_cells) | ||
|
|
There was a problem hiding this comment.
Thanks Asma! How is that different to the original implementation? What does it fix exactly?
The only difference I see is that DiffusionForecastEngine doesn't get self.forecast_engine but initializes directly the ForecastingEngine().
There was a problem hiding this comment.
I like the getattr, which definitely is cleaner. Otherwise, I'd favor to stick to the single if without the else case.
config/diffusion_config.yml
Outdated
There was a problem hiding this comment.
It's clean to have the diffusion config in a separate file. Not sure if this is needed tough at the moment. Maybe @MatKbauer has some thoughts?
There was a problem hiding this comment.
Yes, was thinking so too. Let's leave the config messy as of now during the diffusion forecast engine testing and later revert back to the original defaults (before we actually merge it into develop). @sbAsma, can you hold the changes to the default_config.yml back for now?
There was a problem hiding this comment.
Concerning the separate file for the diffusion engine. This would help us prevent the default_config.yml growing incredibly large. On the other hand, we were concerned that with a nested config structure, we will have hard times to find certain parameters (as we would not be able to search for them with ctrl+F5).
config/diffusion_config.yml
Outdated
There was a problem hiding this comment.
Yes, was thinking so too. Let's leave the config messy as of now during the diffusion forecast engine testing and later revert back to the original defaults (before we actually merge it into develop). @sbAsma, can you hold the changes to the default_config.yml back for now?
src/weathergen/model/model.py
Outdated
| ) | ||
| else: | ||
| self.forecast_engine = ForecastingEngine(cf, self.num_healpix_cells) | ||
|
|
There was a problem hiding this comment.
I like the getattr, which definitely is cleaner. Otherwise, I'd favor to stick to the single if without the else case.
config/diffusion_config.yml
Outdated
There was a problem hiding this comment.
Concerning the separate file for the diffusion engine. This would help us prevent the default_config.yml growing incredibly large. On the other hand, we were concerned that with a nested config structure, we will have hard times to find certain parameters (as we would not be able to search for them with ctrl+F5).
…ssue1279_noise_conditioning
|
@MatKbauer Brought back the previous config ;) |
…com/sbAsma/WeatherGenerator into sbAsma/issue1279_noise_conditioning
tjhunter
changed the title
|
@MatKbauer @sbAsma : Can we merge or close this issue? |
6 participants
Description
changes in
config/default_config.ymlrestored the file to the version in develop and put the diffusion specific changes in a separate config:
config/diffusion_config.ymlchanges in
src/weathergen/model/model.pyadded a code snippet that tests the config param
fe_diffusion_modeland adapts the corresponding forecast engineIssue Number
Checklist before asking for review
./scripts/actions.sh lint./scripts/actions.sh unit-test./scripts/actions.sh integration-testlaunch-slurm.py --time 60