RNO and dissipative loss (MNO) #248
Conversation
Conflicts: neuralop/models/__init__.py neuralop/training/losses.py
neuralop/models/rno.py
Outdated
| if self.domain_padding: | ||
| h = h.unsqueeze(1) # add dim for padding compatibility | ||
| h = self.domain_padding.unpad(h) | ||
| h = h[:,0] # remove extraneous dim |
There was a problem hiding this comment.
Are you referring to squeeze for the last line? In general, if I were to squeeze here, I would be wary of accidentally deleting the out_channels dimension if it is 1.
There was a problem hiding this comment.
Yes, just squeeze that particular dimension, ie. h = h.squeeze(1)
| pred_x = self.layers[i](x, init_hidden_states[i]) | ||
| if i < self.n_layers - 1: | ||
| if self.residual: | ||
| x = x + pred_x |
There was a problem hiding this comment.
I think you need to apply self.layers[i].transform to x in case of output_scaling factor passed by user
There was a problem hiding this comment.
I don't think I have implemented a transform() function for the recurrent layers. Do I need to do this? I thought any output_scale_factor differences were handled by the Fourier blocks under the hood.
There was a problem hiding this comment.
Yes but what happens if self.layers[i] downscales (or upscales) the input? Or does this never happen?
There was a problem hiding this comment.
If I understand correctly, we need transform after each spectral convolution in case we upscale/downscale the input? In recurrent_layers.py, I operate directly on FNO blocks, so I think any changes in resolution are dealt with inside the blocks. What do you think?
There was a problem hiding this comment.
Thank you @mliuschi, it looks much better! I left some comments, mostly just style and clarity.
neuralop/losses/data_losses.py
Outdated
| else: | ||
| self.domain_shape = None | ||
|
|
||
| def sample_uniform_spherical_shell(self, shape: tuple): # shape is (batch_size, dim1, ..., dimk, out_dim) |
There was a problem hiding this comment.
Can we add a docstring here to explain what the shape correspond to, what the function does and how it should be used?
neuralop/losses/data_losses.py
Outdated
| ndim = torch.prod(torch.tensor(shape)) | ||
| inner_radius, outer_radius = self.radii | ||
|
|
||
| samp_radii = torch.zeros(npoints, 1).uniform_(inner_radius, outer_radius) |
There was a problem hiding this comment.
sample_radii? I like to spell out the variables.
neuralop/losses/data_losses.py
Outdated
| domain_shape : int tuple | ||
| Shape of PDE domain (resolution) | ||
| """ | ||
| def __init__(self, data_loss, diss_y_rule, loss_weight: float, diss_radii: tuple, out_dim: int, domain_shape: tuple): |
There was a problem hiding this comment.
Why not explicitly pass the inner and outer radius for clarity instead of a tuple?
Should we call out_dim out_channels? And domain shape, is it just the spatial domain? If so we can say in the docstring, excludes channels.
There was a problem hiding this comment.
As I have it set up, the inner and outer radii are saved as fields in the DissipativeLoss object. What do you think?
neuralop/losses/data_losses.py
Outdated
| """Dissipative loss is used to avoid model blow-up and to encourage the model | ||
| to learn dissipative dynamics (particularly in chaotic systems). The dissipative | ||
| loss samples points uniformly and randomly in a spherical shell around the | ||
| attractor of a chaotic system (at user-specified radii) and encourages |
There was a problem hiding this comment.
Csn you expand on the radii -let's also describe each radius individually?
neuralop/losses/data_losses.py
Outdated
| loss samples points uniformly and randomly in a spherical shell around the | ||
| attractor of a chaotic system (at user-specified radii) and encourages | ||
| that a learned model's predictions on this shell scale inwards by some | ||
| user-specified dissipative rule. |
There was a problem hiding this comment.
Can you have some higher, less technical level intuition here? Readers can refer to the paper for details.
neuralop/losses/data_losses.py
Outdated
| samp_radii = torch.zeros(npoints, 1).uniform_(inner_radius, outer_radius) | ||
| vec = torch.randn(npoints, ndim) # ref: https://mathworld.wolfram.com/SpherePointPicking.html | ||
| vec = torch.nn.functional.normalize(vec, p=2.0, dim=1) | ||
| samp_radii = torch.broadcast_to(samp_radii, vec.shape) |
There was a problem hiding this comment.
Same, try not to change the semantic meaning of variables half way through as it makes the code quite hard to read.
neuralop/losses/data_losses.py
Outdated
|
|
||
| return (samp_radii * vec).reshape((npoints, *shape)) | ||
|
|
||
| def __call__(self, x, y, x_diss, pred_diss, **kwargs): |
There was a problem hiding this comment.
Add a short docstring as the arguments are non-standard.
Usually our losses take as input y_pred, y.
There was a problem hiding this comment.
Done! I have also added a "Usage example" in the main docstring for the class as well
| pred_x = self.layers[i](x, init_hidden_states[i]) | ||
| if i < self.n_layers - 1: | ||
| if self.residual: | ||
| x = x + pred_x |
There was a problem hiding this comment.
Yes but what happens if self.layers[i] downscales (or upscales) the input? Or does this never happen?
neuralop/models/rno.py
Outdated
| if self.domain_padding: | ||
| h = h.unsqueeze(1) # add dim for padding compatibility | ||
| h = self.domain_padding.unpad(h) | ||
| h = h[:,0] # remove extraneous dim |
There was a problem hiding this comment.
Yes, just squeeze that particular dimension, ie. h = h.squeeze(1)
No description provided.