Add CNN predictor#645
Conversation
Co-authored-by: Codex <codex@openai.com>
Co-authored-by: Codex <codex@openai.com>
|
@joeloskarsson please check it .I will move forward after it |
Co-authored-by: Codex <codex@openai.com>
|
Please add the full PR template. Is the CNN the backbone from https://arxiv.org/abs/2507.05658? Would be good to add a reference. I would maybe add just a U-net as a first CNN architecture (to figure out the details that comes with requiring a regular grid), but I am not opposed to this one. |
|
@joeloskarsson Yes this is ResHRRR Inspired i used the exact blocks used in the paper but not the full implementation . |
Co-authored-by: Codex <codex@openai.com>
|
@joeloskarsson yes it is done .I have not addded the unet backbone .If you want i can add |
Sir-Sloth-The-Lazy
left a comment
Sir-Sloth-The-Lazy
left a comment
There was a problem hiding this comment.
Hi, sorry for the empty review before, I this is really nice, just some issue that I found and thought might bring it to your attention. Thank for your time !
| padding_mode: str = "zeros", | ||
| ): | ||
| """Initialize a shape-preserving residual convolution block.""" | ||
| super().__init__() | ||
| if channels <= 0: | ||
| raise ValueError("channels must be positive") | ||
| if kernel_size <= 0 or kernel_size % 2 == 0: | ||
| raise ValueError("kernel_size must be a positive odd integer") | ||
|
|
||
| padding = kernel_size // 2 | ||
| self.conv1 = nn.Conv2d( | ||
| channels, | ||
| channels, | ||
| kernel_size=kernel_size, | ||
| padding=padding, | ||
| padding_mode=padding_mode, | ||
| ) | ||
| self.act = nn.SiLU() | ||
| self.conv2 = nn.Conv2d( | ||
| channels, | ||
| channels, | ||
| kernel_size=kernel_size, | ||
| padding=padding, | ||
| padding_mode=padding_mode, |
There was a problem hiding this comment.
padding_mode gets passed straight into nn.Conv2d here with nothing checking it against the actual grid shape. PyTorch's reflect mode needs padding < input_size on the padded axis, and circular needs the padding to not wrap around more than once. A grid with a spatial dimension of 1 will crash on the first forward pass here even with the default cnn_kernel_size=3 and reflect, and circular will crash on bigger grids too once the kernel gets large enough. Since --cnn_padding_mode is just an open CLI choice right now, nothing stops someone from hitting this. Might be worth validating grid shape against kernel size and padding mode at construction time, or at least documenting the constraint.
| ) | ||
| grid_features = node_to_grid(grid_features, self.grid_shape) | ||
|
|
||
| context = forcing.mean(dim=1) if self.cnn_film else None |
There was a problem hiding this comment.
This averages forcing over every grid node, so every cell ends up with the exact same FiLM conditioning no matter what the local forcing looks like (boundary vs interior, etc). Not a crash, just seems like it throws away the spatial signal a regional CNN would otherwise want to use per pixel. Was this a deliberate choice (matching how HRRRCast conditions globally), or was per node conditioning the intent?
| da_state_stats = datastore.get_standardization_dataarray("state") | ||
| self.register_buffer( | ||
| "diff_mean", | ||
| torch.tensor( | ||
| da_state_stats.state_diff_mean_standardized.values, | ||
| dtype=torch.float32, | ||
| ), | ||
| persistent=False, | ||
| ) | ||
| self.register_buffer( | ||
| "diff_std", | ||
| torch.tensor( | ||
| da_state_stats.state_diff_std_standardized.values, | ||
| dtype=torch.float32, | ||
| ), | ||
| persistent=False, | ||
| ) |
There was a problem hiding this comment.
This block is basically identical to what's already in BaseGraphModel.__init__ (neural_lam/models/step_predictors/graph/base.py, around L76 to L89). Since any delta predicting StepPredictor needs this rescale logic, seems like a good candidate to move up into the shared StepPredictor.__init__ next to state_mean/state_std rather than copying it into every subclass.
|
Okay I will look into it |
Co-authored-by: Codex <codex@openai.com>
|
@joeloskarsson i need yor review .This will be my first PR for the a issue who is being assigned to me so i am excited for your comments |
4 participants
Describe your changes
Adds a regular-grid CNN predictor path for Neural-LAM.
This PR adds reusable node/grid tensor transforms, ResHRRR-inspired CNN layers, a
CNNPredictorStepPredictor, CLI construction for--model cnn_predictor, checkpoint reconstruction support, and focused tests.The CNN backbone is inspired by the ResHRRR model from HRRRCast: https://arxiv.org/abs/2507.05658.
It includes residual CNN blocks, Squeeze-and-Excitation, and optional FiLM conditioning, but it is not a full HRRRCast reproduction.
Multi-lead training, DDIM/diffusion sampling, and a U-Net backbone are left out of this PR.
No new runtime dependency is required.
Issue Link
refs #590
Type of change
Checklist before requesting a review
pullwith--rebaseoption if possible).Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
Author checklist after completed review
reflecting type of change (add section where missing):
Checklist for assignee