Implement snapshotting for the acoustic wave equation #2474
Conversation
| # Build operator equations | ||
| equations = eqn + src_term + rec_term | ||
|
|
||
| if factor: |
There was a problem hiding this comment.
This needs to be wrapped into a utility function as it's duplicated below
| nsnaps = (geometry.nt + factor - 1) // factor | ||
| time_subsampled = ConditionalDimension( | ||
| 't_sub', parent=model.grid.time_dim, factor=factor) | ||
| usnaps = TimeFunction(name='usnaps', grid=model.grid, |
There was a problem hiding this comment.
You still have u with full time saved line 135 you can't have both
| name='Forward', **kwargs) | ||
| op = Operator(equations, subs=model.spacing_map, name='Forward', **kwargs) | ||
| if usnaps is not None: | ||
| return op, usnaps |
There was a problem hiding this comment.
No the operator build cannot return objects like that. This is an abstract operator with placeholders that might not be correct for runtime.
|
|
||
| if factor is not None: | ||
| # Condition to apply gradient update only at snapshot times | ||
| condition = Eq(time % factor, 0) |
There was a problem hiding this comment.
No you don't need that usnap already contains the conditon
| u = TimeFunction(name='u', grid=model.grid, | ||
| save=geometry.nt if save else None, | ||
| time_order=2, space_order=space_order) | ||
| if kernel == 'OT2': |
There was a problem hiding this comment.
Unnecessary duplicate, u contains the information you should not need separate cases for gradient_update
| @@ -1,6 +1,6 @@ | |||
| from devito import Function, TimeFunction, DevitoCheckpoint, CheckpointOperator, Revolver | |||
| from devito.tools import memoized_meth | |||
| from examples.seismic.acoustic.operators import ( | |||
| from devitofwi.devito.acoustic.operators import ( | |||
| @@ -108,12 +111,24 @@ def forward(self, src=None, rec=None, u=None, model=None, save=None, **kwargs): | |||
| model = model or self.model | |||
| # Pick vp from model unless explicitly provided | |||
| kwargs.update(model.physical_params(**kwargs)) | |||
| # Get the operator | |||
| op_fwd = self.op_fwd(save=save, factor=factor) | |||
| # Prepare parameters for operator apply | |||
There was a problem hiding this comment.
Don't know what this is for.
| dt=kwargs.pop('dt', self.dt), **kwargs) | ||
|
|
||
| return rec, u, summary | ||
| if factor: |
There was a problem hiding this comment.
No, usnap needs to be create here like u then passed as argument
| op_args['usnaps'] = usnaps | ||
| summary = op.apply(**op_args) | ||
|
|
||
| else: |
There was a problem hiding this comment.
shouldn't need if else only kwargs
| @@ -209,8 +236,17 @@ def jacobian_adjoint(self, rec, u, src=None, v=None, grad=None, model=None, | |||
| wrp.apply_forward() | |||
| summary = wrp.apply_reverse() | |||
| else: | |||
| summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, dt=dt, | |||
| **kwargs) | |||
| if factor is not None: | |||
There was a problem hiding this comment.
Again, not needed, input u should contain all metada needed
There was a problem hiding this comment.
Hi Mathias, thank you for your feedback. I have reviewed and cleaned the code to the best of my understanding. I have included a notebook to compare computing the FWI gradient with and without snapshotting and two scripts to calculate the memory usage of both methods. After updating the code, the memory usage for calculating the gradient with snapshotting is more than twice that of the older code version. This reduced memory usage (I guess) because I was passing 'usnaps' with the operator (which is not good practice). I am wondering, is it possible to improve the code more to reduce the memory usage?
…apshot_acoustic
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Implement snapshotting to save snapshots of the forward wavefield used to compute the gradient to reduce memory usage.