import warnings import torch import torch.nn as nn from .odeint import SOLVERS, odeint from .misc import _check_inputs, _flat_to_shape, _mixed_norm, _all_callback_names, _all_adjoint_callback_names class OdeintAdjointMethod(torch.autograd.Function): @staticmethod def forward(ctx, shapes, func, y0, t, rtol, atol, method, options, event_fn, adjoint_rtol, adjoint_atol, adjoint_method, adjoint_options, t_requires_grad, *adjoint_params): ctx.shapes = shapes ctx.func = func ctx.adjoint_rtol = adjoint_rtol ctx.adjoint_atol = adjoint_atol ctx.adjoint_method = adjoint_method ctx.adjoint_options = adjoint_options ctx.t_requires_grad = t_requires_grad ctx.event_mode = event_fn is not None with torch.no_grad(): ans = odeint(func, y0, t, rtol=rtol, atol=atol, method=method, options=options, event_fn=event_fn) if event_fn is None: y = ans ctx.save_for_backward(t, y, *adjoint_params) else: event_t, y = ans ctx.save_for_backward(t, y, event_t, *adjoint_params) return ans @staticmethod def backward(ctx, *grad_y): with torch.no_grad(): func = ctx.func adjoint_rtol = ctx.adjoint_rtol adjoint_atol = ctx.adjoint_atol adjoint_method = ctx.adjoint_method adjoint_options = ctx.adjoint_options t_requires_grad = ctx.t_requires_grad # Backprop as if integrating up to event time. # Does NOT backpropagate through the event time. event_mode = ctx.event_mode if event_mode: t, y, event_t, *adjoint_params = ctx.saved_tensors _t = t t = torch.cat([t[0].reshape(-1), event_t.reshape(-1)]) grad_y = grad_y[1] else: t, y, *adjoint_params = ctx.saved_tensors grad_y = grad_y[0] adjoint_params = tuple(adjoint_params) ################################## # Set up initial state # ################################## # [-1] because y and grad_y are both of shape (len(t), *y0.shape) aug_state = [torch.zeros((), dtype=y.dtype, device=y.device), y[-1], grad_y[-1]] # vjp_t, y, vjp_y aug_state.extend([torch.zeros_like(param) for param in adjoint_params]) # vjp_params ################################## # Set up backward ODE func # ################################## # TODO: use a nn.Module and call odeint_adjoint to implement higher order derivatives. def augmented_dynamics(t, y_aug): # Dynamics of the original system augmented with # the adjoint wrt y, and an integrator wrt t and args. y = y_aug[1] adj_y = y_aug[2] # ignore gradients wrt time and parameters with torch.enable_grad(): t_ = t.detach() t = t_.requires_grad_(True) y = y.detach().requires_grad_(True) # If using an adaptive solver we don't want to waste time resolving dL/dt unless we need it (which # doesn't necessarily even exist if there is piecewise structure in time), so turning off gradients # wrt t here means we won't compute that if we don't need it. func_eval = func(t if t_requires_grad else t_, y) # Workaround for PyTorch bug #39784 _t = torch.as_strided(t, (), ()) # noqa _y = torch.as_strided(y, (), ()) # noqa _params = tuple(torch.as_strided(param, (), ()) for param in adjoint_params) # noqa vjp_t, vjp_y, *vjp_params = torch.autograd.grad( func_eval, (t, y) + adjoint_params, -adj_y, allow_unused=True, retain_graph=True ) # autograd.grad returns None if no gradient, set to zero. vjp_t = torch.zeros_like(t) if vjp_t is None else vjp_t vjp_y = torch.zeros_like(y) if vjp_y is None else vjp_y vjp_params = [torch.zeros_like(param) if vjp_param is None else vjp_param for param, vjp_param in zip(adjoint_params, vjp_params)] return (vjp_t, func_eval, vjp_y, *vjp_params) # Add adjoint callbacks for callback_name, adjoint_callback_name in zip(_all_callback_names, _all_adjoint_callback_names): try: callback = getattr(func, adjoint_callback_name) except AttributeError: pass else: setattr(augmented_dynamics, callback_name, callback) ################################## # Solve adjoint ODE # ################################## if t_requires_grad: time_vjps = torch.empty(len(t), dtype=t.dtype, device=t.device) else: time_vjps = None for i in range(len(t) - 1, 0, -1): if t_requires_grad: # Compute the effect of moving the current time measurement point. # We don't compute this unless we need to, to save some computation. func_eval = func(t[i], y[i]) dLd_cur_t = func_eval.reshape(-1).dot(grad_y[i].reshape(-1)) aug_state[0] -= dLd_cur_t time_vjps[i] = dLd_cur_t # Run the augmented system backwards in time. aug_state = odeint( augmented_dynamics, tuple(aug_state), t[i - 1:i + 1].flip(0), rtol=adjoint_rtol, atol=adjoint_atol, method=adjoint_method, options=adjoint_options ) aug_state = [a[1] for a in aug_state] # extract just the t[i - 1] value aug_state[1] = y[i - 1] # update to use our forward-pass estimate of the state aug_state[2] += grad_y[i - 1] # update any gradients wrt state at this time point if t_requires_grad: time_vjps[0] = aug_state[0] # Only compute gradient wrt initial time when in event handling mode. if event_mode and t_requires_grad: time_vjps = torch.cat([time_vjps[0].reshape(-1), torch.zeros_like(_t[1:])]) adj_y = aug_state[2] adj_params = aug_state[3:] return (None, None, adj_y, time_vjps, None, None, None, None, None, None, None, None, None, None, *adj_params) def odeint_adjoint(func, y0, t, *, rtol=1e-7, atol=1e-9, method=None, options=None, event_fn=None, adjoint_rtol=None, adjoint_atol=None, adjoint_method=None, adjoint_options=None, adjoint_params=None): # We need this in order to access the variables inside this module, # since we have no other way of getting variables along the execution path. if adjoint_params is None and not isinstance(func, nn.Module): raise ValueError('func must be an instance of nn.Module to specify the adjoint parameters; alternatively they ' 'can be specified explicitly via the `adjoint_params` argument. If there are no parameters ' 'then it is allowable to set `adjoint_params=()`.') # Must come before _check_inputs as we don't want to use normalised input (in particular any changes to options) if adjoint_rtol is None: adjoint_rtol = rtol if adjoint_atol is None: adjoint_atol = atol if adjoint_method is None: adjoint_method = method if adjoint_method != method and options is not None and adjoint_options is None: raise ValueError("If `adjoint_method != method` then we cannot infer `adjoint_options` from `options`. So as " "`options` has been passed then `adjoint_options` must be passed as well.") if adjoint_options is None: adjoint_options = {k: v for k, v in options.items() if k != "norm"} if options is not None else {} else: # Avoid in-place modifying a user-specified dict. adjoint_options = adjoint_options.copy() if adjoint_params is None: adjoint_params = tuple(find_parameters(func)) else: adjoint_params = tuple(adjoint_params) # in case adjoint_params is a generator. # Filter params that don't require gradients. oldlen_ = len(adjoint_params) adjoint_params = tuple(p for p in adjoint_params if p.requires_grad) if len(adjoint_params) != oldlen_: # Some params were excluded. # Issue a warning if a user-specified norm is specified. if 'norm' in adjoint_options and callable(adjoint_options['norm']): warnings.warn("An adjoint parameter was passed without requiring gradient. For efficiency this will be " "excluded from the adjoint pass, and will not appear as a tensor in the adjoint norm.") # Convert to flattened state. shapes, func, y0, t, rtol, atol, method, options, event_fn, decreasing_time = _check_inputs(func, y0, t, rtol, atol, method, options, event_fn, SOLVERS) # Handle the adjoint norm function. state_norm = options["norm"] handle_adjoint_norm_(adjoint_options, shapes, state_norm) ans = OdeintAdjointMethod.apply(shapes, func, y0, t, rtol, atol, method, options, event_fn, adjoint_rtol, adjoint_atol, adjoint_method, adjoint_options, t.requires_grad, *adjoint_params) if event_fn is None: solution = ans else: event_t, solution = ans event_t = event_t.to(t) if decreasing_time: event_t = -event_t if shapes is not None: solution = _flat_to_shape(solution, (len(t),), shapes) if event_fn is None: return solution else: return event_t, solution def find_parameters(module): assert isinstance(module, nn.Module) # If called within DataParallel, parameters won't appear in module.parameters(). if getattr(module, '_is_replica', False): def find_tensor_attributes(module): tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v) and v.requires_grad] return tuples gen = module._named_members(get_members_fn=find_tensor_attributes) return [param for _, param in gen] else: return list(module.parameters()) def handle_adjoint_norm_(adjoint_options, shapes, state_norm): """In-place modifies the adjoint options to choose or wrap the norm function.""" # This is the default adjoint norm on the backward pass: a mixed norm over the tuple of inputs. def default_adjoint_norm(tensor_tuple): t, y, adj_y, *adj_params = tensor_tuple # (If the state is actually a flattened tuple then this will be unpacked again in state_norm.) return max(t.abs(), state_norm(y), state_norm(adj_y), _mixed_norm(adj_params)) if "norm" not in adjoint_options: # `adjoint_options` was not explicitly specified by the user. Use the default norm. adjoint_options["norm"] = default_adjoint_norm else: # `adjoint_options` was explicitly specified by the user... try: adjoint_norm = adjoint_options['norm'] except KeyError: # ...but they did not specify the norm argument. Back to plan A: use the default norm. adjoint_options['norm'] = default_adjoint_norm else: # ...and they did specify the norm argument. if adjoint_norm == 'seminorm': # They told us they want to use seminorms. Slight modification to plan A: use the default norm, # but ignore the parameter state def adjoint_seminorm(tensor_tuple): t, y, adj_y, *adj_params = tensor_tuple # (If the state is actually a flattened tuple then this will be unpacked again in state_norm.) return max(t.abs(), state_norm(y), state_norm(adj_y)) adjoint_options['norm'] = adjoint_seminorm else: # And they're using their own custom norm. if shapes is None: # The state on the forward pass was a tensor, not a tuple. We don't need to do anything, they're # already going to get given the full adjoint state as (t, y, adj_y, adj_params) pass # this branch included for clarity else: # This is the bit that is tuple/tensor abstraction-breaking, because the odeint machinery # doesn't know about the tupled nature of the forward state. We need to tell the user's adjoint # norm about that ourselves. def _adjoint_norm(tensor_tuple): t, y, adj_y, *adj_params = tensor_tuple y = _flat_to_shape(y, (), shapes) adj_y = _flat_to_shape(adj_y, (), shapes) return adjoint_norm((t, *y, *adj_y, *adj_params)) adjoint_options['norm'] = _adjoint_norm