Split preactivation from FNOBlock.forward()
#214
Conversation
Create a separate method `FNOBlock.forward_with_preactivation()` where it is assumed `preactivation=True` for clarity of code flow.
In multiple places, we were checking `self.output_scaling_factors` and calling `resample` like boilerplate. Refactor this check and call (with a no-op return if `scaling_factors=None`). Additionally, refactor out `output_scaling_factors` validation, which was being duplicated across 4 files.
neuralop/utils.py
Outdated
|
|
||
| def validate_output_scaling_factor_1d( | ||
| output_scaling_factor: Optional[Union[Number, List[Number]]], | ||
| order: int, |
There was a problem hiding this comment.
Why do you call them 1d and 2d? They are both N-D, where order = n_dim = N.
Given this I think we should merge them and just have n_layers=1 by default, what do you think?
There was a problem hiding this comment.
Why do you call them 1d and 2d?
That's just based on their return value. validate_*_2d returns a list of shapes, while validate_*_1d returns just a shape.
I think we should merge them and just have n_layers=1 by default
That sounds fine, as long as a singleton "list" of shapes doesn't confuse any of the models.
There was a problem hiding this comment.
I see - how about we call them validate_output_scaling_factor_single_layer and validate_output_scaling_factor_multiple_layers?
A mouthful and a pain to write but self-explanatory and never used by end users most likely?
There was a problem hiding this comment.
Ideally we should never have to do both, makes me think there is a deficiency in the code - we should either check at model level or layer level, probably not both
There was a problem hiding this comment.
I made an allowance for non-positive values of n_layers to return a single layer instead of a list of layers (for a whole module). This is only used in DomainPadding.
|
|
||
| if init_std == "auto": | ||
| init_std = 1 / (in_channels * out_channels) | ||
| else: | ||
| init_std = 0.02 | ||
| init_std = init_std |
There was a problem hiding this comment.
Thanks, I kept forgetting to update this!
Also, run `black uno.py tests.test_uno.py`
| and (self.mlp is not None) | ||
| or (index < (self.n_layers - index)) | ||
| ): | ||
| if (self.mlp is not None) or (index < (self.n_layers - index)): | ||
| x = self.non_linearity(x) |
There was a problem hiding this comment.
The order of operations (i.e. and , or) in this expression confused me, but I assumed this should only be applied in the preactivation=False case.
If n_layers<=0, return a list, not a list of lists
neuralop/layers/fno_block.py
Outdated
| return x | ||
|
|
||
| def resample(self, x, index, output_shape): | ||
| """Resamples input if scaling factors are available for this block.""" | ||
| if self.output_scaling_factor is None: |
There was a problem hiding this comment.
Shouldn't you check both self.output_scaling_factor and output_shape vs x.shape?
There was a problem hiding this comment.
Good point. I've updated FNOBlock.resample() to defer first to arg output_shape to be passed to resample. AFAICT, when output_shape is not None in resample.py:resample() the other args are ignored (I passed dummy args res_scale=1, axis=None, so it should be fine).
|
This is great, thanks for all the changes @m4e7! Merging. |
Mainly, refactor
FNOBlock.forward()for readability:FNOBlock.forward_with_preactivation()where it is assumedpreactivation=Truefor clarity of code flow. Likewise, assumepreactivation=Falsein the new methodFNOBlock.forward_with_postactivation(). The old methodforward()now dispatches to one of these.self.output_scaling_factorsand callingresamplelike boilerplate. Refactor this check and call (with a no-op return ifscaling_factors is None and output_shape is None).output_scaling_factorsvalidation, which was being duplicated across 4 files.Reviewer(s): @JeanKossaifi @zongyi-li