Skip to main content

ASDL: Automatic Second-order Differentiation (for Fisher, Gradient covariance, Hessian, Jacobian, and Kernel) Library

Project description

ASDL: Automatic Second-order Differentiation Library

ASDL is an extension library of PyTorch to easily perform gradient preconditioning using second-order information (e.g., Hessian, Fisher information) for deep neural networks.

ASDL provides various implementations and a unified interface (GradientMaker) for gradient preconditioning for deep neural networks. For example, to train your model with gradient preconditioning by K-FAC algorithm, you can replace a <Standard> gradient calculation procedure (i.e., a forward pass followed by a backward pass) with one by <ASDL> with KfacGradientMaker like the following:

from asdl.precondition import PreconditioningConfig, KfacGradientMaker

# Initialize model
model = Net()

# Initialize optimizer (SGD is recommended)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# Initialize KfacGradientMaker
config = PreconditioningConfig(data_size=batch_size, damping=0.01)
gm = KfacGradientMaker(model, config)

# Training loop
for x, t in data_loader:
  optimizer.zero_grad()
  
  # <Standard> (gradient calculation)
  # y = model(x)
  # loss = loss_fn(y, t)
  # loss.backward()

  # <ASDL> ('preconditioned' gradient calculation)
  dummy_y = gm.setup_model_call(model, x)
  gm.setup_loss_call(loss_fn, dummy_y, t)
  y, loss = gm.forward_and_backward()

  optimizer.step()

You can apply a different gradient preconditioning algorithm by replacing gm with another XXXGradientMaker(model, config) (XXX: algorithm name, e.g., ShampooGradientMaker for Shampoo algorithm) with the same interface. This enables a flexible switching/comparison of a range of gradient preconditioning algorithms.

Installation

You can install the latest version of ASDL by running:

$ pip install asdfghjkl

ASDL is tested with Python 3.7 and is compatible with PyTorch 1.13.

Example

The training script for training MLPs, CNNs, or ResNets using varous types of gradient preconditionig methods (which reproduces the results in the ASDL paper).

Resource

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

asdfghjkl-0.1a4.tar.gz (75.7 kB view details)

Uploaded Source

Built Distribution

asdfghjkl-0.1a4-py3-none-any.whl (89.7 kB view details)

Uploaded Python 3

File details

Details for the file asdfghjkl-0.1a4.tar.gz.

File metadata

  • Download URL: asdfghjkl-0.1a4.tar.gz
  • Upload date:
  • Size: 75.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.9.4

File hashes

Hashes for asdfghjkl-0.1a4.tar.gz
Algorithm Hash digest
SHA256 a934411a0ffdee6fcdccb19672196498ea6a8e55e3e67abbe67200c84b46ddee
MD5 b936ad5784c919929072f5cb45f7e56a
BLAKE2b-256 c573c81489fe1ed8b12099c8e034e8e623374016e5e6db2010c38d8f386acc2d

See more details on using hashes here.

File details

Details for the file asdfghjkl-0.1a4-py3-none-any.whl.

File metadata

  • Download URL: asdfghjkl-0.1a4-py3-none-any.whl
  • Upload date:
  • Size: 89.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.9.4

File hashes

Hashes for asdfghjkl-0.1a4-py3-none-any.whl
Algorithm Hash digest
SHA256 9cedec38494d7907075941d03d018d8c2b64047e6c41713a49f30e3790816c51
MD5 5cdd6f457d28249dea0c576c76e3ccf8
BLAKE2b-256 329a7548968acfacb162c88dd9c2247a54f36136221ebb13d457016b28d47523

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page