Speed up Jacobian matrix calculation in Pytorch

Hi all
I use ‘torch.autograd.functional.jacobian(f,x)’ to calculate the partial derivatives of f with respect to x but when the dimension of f is incremented the time of calculation increases.
does anyone know any solution to speed up the jacobian calculation in PyTorch?

Thanks in advance.

You might want to have a look at FuncTorch

It has an issue…
Successfully built functorch
Installing collected packages: functorch
Successfully installed functorch-0.2.0a0+a2d2f99
but when I run my code
from functorch import vmap, grad
ModuleNotFoundError: No module named ‘functorch’

Did you install it within the current environment? For example, running pytorch in pip but installing functorch with conda?

I installed the package in my current environment using pip command as had been described in the installation guide. But, the name of my current environment is test.

If you’re using Linux, you can check what python3 executable you’re using with which python, if you installed FuncTorch you most likely made a new environment (as is stated in the install guide) and installed it in that. Also check you didn’t install pytorch with conda when you installed! Then if that’s all ok, manually check your environemnt and see where pytorch is installed and where functorch is installed?

It doesn’t work on my local machine. While installing the functorch package some errors occur.

Did you follow the instruction of the google colab? When you install functorch in a new environment you’ll need to ensure you’ve PyTorch (among other dependencies) reinstalled for the new environment too.

Yes, I followed them without any problem but when my google drive was mounted to load my code which used the’ functorch’ package the error of there is no package by the name of the functorch was shown.

Try installing after mounting your google drive? It sounds like you’re using different environments.

As I checked google colab guid I don’t need to create an environment. Am I right?

I mounted my google drive at the first then the functorch package was installed. Finally, I tried to run my code by the colab unfortunately an error has occurred.

tcmalloc: large alloc 140349731364864 bytes == (nil) @ 0x7fa5b927e887 0x7fa5b7b74c29 0x7fa5b7b75afb 0x7fa5b7b7620c 0x7fa4f64e6418 0x7fa4f64ea2d8 0x7fa4aee1bd1d 0x7fa4aee151b1 0x7fa5b94c98d3 0x7fa5b94ce39f 0x7fa5b81ac1ef 0x7fa5b94cd96a 0x7fa5b8e26f96 0x7fa5b81ac1ef 0x7fa5b81ac27f 0x7fa5b8e27745 0x7fa5b8e27051 0x7fa5ab7a3a16 0x55d25a847317 0x55d25a847120 0x55d25a8baf33 0x55d25a8b666e 0x55d25a849151 0x55d25a88a7b9 0x55d25a8876d4 0x55d25a847b32 0x55d25a8bba6d 0x55d25a8489da 0x55d25a8bb2c0 0x55d25a8b602f 0x55d25a8b5d43
terminate called after throwing an instance of ‘std::bad_alloc’
what(): std::bad_alloc

Do you have any advice?

Looks like you’re tring to allocate a lot of RAM. It’s most likely a common error on google-drive and I’d check on stackexchange to see if it’s been noted there!

I’d check it by using batch 8! it got the error again!!!