Closed jharris1679 closed 1 year ago
Hi!
You can enable grad during validation. You need to pass Trainer(inference_mode=False)
and then you are free to use the context manager.
Duplicate of https://github.com/Lightning-AI/lightning/issues/13948
Oh fantastic, I'll check this out, thank you!
I've tried this but still seeing RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
during validation steps, but I can get gradients successfully during a training step. Is there something else I'm missing?
You will need to set tensor.requires_grad_()
for the tensors you want to compute grads for.
Thanks for such a quick response! In my training step, I don't need to do that for the net output, and the loss function is working well. However in my validation_step, if I specify requires grad like so
output = self.net(x)
output.requires_grad_(requires_grad=True)
and then take the grads like so
grads = torch.autograd.grad(
output,
[x],
grad_outputs=torch.ones_like(output),
create_graph=True,
retain_graph=True
)
it gives
RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behaviour.
which is the same error I would get if I didnt set inference_mode=False
. If I set allow_unused=True
in torch.autograd.grad
then grads are None
. It seems as though my Trainer(inference_mode=False)
is maybe not doing the trick? Should I perhaps open another issue with a steps to reproduce?
Thanks again for your eyes on this!
Ahh in addition to .requires_grad_()
I had to put everything in my validation step under with torch.enable_grad():
, and now it works! Thanks again.
Have you tried using functorch to compute the partials ? It's much easier to use.
I have an example script here. https://github.com/GeoffNN/deeponet-fno/blob/main/src/burgers/pytorch_deeponet.py
🚀 Feature
Enable gradient calculations during evaluation loops.
Motivation
Some loss functions require the gradients of the outputs with respect to the inputs. For example, a physics informed neural network uses these gradients in a differential equation as its loss function.
Pitch
Add a
set_grad_enabled
flag to validation step to make the following possible for learning the cosine function, for example:Alternatives
I tried simply adding
with torch.set_grad_enabled(True)
to my code but of course that didn't work.Additional context
Physics Informed Neural Networks are a recently introduced architecture for learning differential equations by embedding them in the loss function of a neural network. The key innovation involves repurposing auto-differentiation machinery to obtain derivatives of network outputs wrt its inputs, then plugging those into the residual form of a differential equation. The network learns accurate derivatives by minimizing this residual.
PINNs paper: https://faculty.sites.iastate.edu/hliu/files/inline-files/PINN_RPK_2019_1.pdf
If you enjoy Lightning, check out our other projects! âš¡
Metrics: Machine learning metrics for distributed, scalable PyTorch applications.
Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.
Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.
Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.
Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging PyTorch Lightning, Transformers, and Hydra.