Use new Jacobian API's in Torch

Created by: gcarleo

PyTorch 1.5 now supports APIs for the Jacobian

We would need to use these API's in our machine.Torch in place of the loop we implemented there

added help wanted v3.0 labels

changed milestone to %Version 3.0

Created by: ChenAo-Phys

Hello. I was also working on the implementation of Jacobian in these days. The new pytorch API for Jacobian seems not a good choice, because its source code shows that this algorithm is still based on for loops. There is a package named backpack which implements the Jacobian in pytorch. Here is the paper https://arxiv.org/pdf/1912.10985.pdf and their website http://backpack.pt/

In fig.3 of that paper they show that the calculation of jacobian is as fast as a single backward pass, and of course much faster than for loops.

The package can be used like this

import torch
from torch import nn
from backpack import backpack, extend
from backpack.extensions import BatchGrad

model = nn.Sequential(
    nn.Conv2d(1, 1, 3),
    nn.ReLU(),
    nn.Conv2d(1, 1, 3),
    nn.ReLU(),
    nn.AvgPool2d(6),
    nn.Flatten(0, -1)
)

model = extend(model)
s = torch.randn(2,1,10,10)
psi = model(s)
psi_sum = psi.sum()
with backpack(BatchGrad()):
    psi_sum.backward()
for param in model.parameters():
    print(param.grad_batch)

Created by: gcarleo

Thank you @ChenAo-Phys indeed the current implementation of the jacobian is not very efficient in pytorch, and most likely it will have the same performance of our already implemented for loop...

Would you be willing to do a PR to include this in our Torch machine?

Created by: ChenAo-Phys

Yes I can add this into Torch machine. I did a test today. The calculation of Jacobian in backpack is really efficient and consumes much less time compared with other parts like iterative solver.

Created by: gcarleo

Amazing @ChenAo-Phys ! Yes please do a PR. I would suggest that you do this PR to the master branch and/or to the v3.0 branch

Created by: gcarleo

@ChenAo-Phys just an implementation note: I saw that backpack does not support many layers (essentially it supports only 2 types of layers...) so I wouldn't make this the default setting to compute der_log in our machine, but maybe some opt-in solution for the user

Created by: ChenAo-Phys

@gcarleo I'm sorry that I only tested with Conv2d and didn't notice this problem. I agree that backpack should not be the default method. Or maybe we can judge whether to use backpack automatically, like

use_backpack = True
for module in model.children():
    if not isinstance(module, (nn.Conv2d, nn.Linear)) \
       and sum(p.numel() for p in module.parameters()):
        use_backpack = False
        break

Created by: gcarleo

yes that looks good! I would also maybe add a check that backpack is importable to decide whether to use it or not

Created by: ChenAo-Phys

@gcarleo Based on the idea of backpack, I improve it to accommodate most pytorch layers including nn.Linear, nn.Conv1d/2d/3d, ... https://github.com/ChenAo-Phys/pytorch-Jacobian This may be a better solution for the current problem.

Created by: gcarleo

Very nice @ChenAo-Phys , (and I guess the people developing backpack might be also interested in integrating this in their code). It would be good if you could do a PR to netket at some point (again, I would do this as an opt-in feature, i.e. use this fast jacobian only if your library is found installed? )

assigned to @filippo.vicentini

Created by: ChenAo-Phys

@gcarleo The library is no more than 200 lines so I think it's better to add it directly into netket/machine/torch.py to avoid the trouble of installing my library for users. The problem is that sometimes the network is not written in a standard form to make my method work. I will do an automatic check and deactivate my method if that happens.

Created by: gcarleo

Yes, that sounds like a good strategy thanks!

Created by: ChenAo-Phys

@gcarleo I ignored a problem in my method - it's too memory-consuming. It will probably fail in PC if one wants to calculate the Jacobian of a network with more than 3000 parameters. It's really a pity. Now I think it's still better to use backpack if it's checked to be importable.

Created by: gcarleo

Closing this because in 3.0 we are dropping support for PyTorch

closed