BackPACK

"""
Compute the gradient with Pytorch

"""
from torch.nn import CrossEntropyLoss, Linear
from backpack.utils.examples import load_one_batch_mnist


X, y = load_one_batch_mnist(flat=True)
model = Linear(784, 10)
lossfunc = CrossEntropyLoss()
loss = lossfunc(model(X), y)


loss.backward()

for param in model.parameters():
    print(param.grad)

"""
Compute the gradient with Pytorch
and the variance with BackPACK
"""
from torch.nn import CrossEntropyLoss, Linear
from backpack.utils.examples import load_one_batch_mnist
from backpack import extend, backpack
from backpack.extensions import Variance

X, y = load_one_batch_mnist(flat=True)
model = extend(Linear(784, 10))
lossfunc = extend(CrossEntropyLoss())
loss = lossfunc(model(X), y)

with backpack(Variance()):
    loss.backward()

for param in model.parameters():
    print(param.grad)
    print(param.variance)

"""
Compute the gradient with Pytorch
and the second moment with BackPACK
"""
from torch.nn import CrossEntropyLoss, Linear
from backpack.utils.examples import load_one_batch_mnist
from backpack import extend, backpack
from backpack.extensions import SumGradSquared

X, y = load_one_batch_mnist(flat=True)
model = extend(Linear(784, 10))
lossfunc = extend(CrossEntropyLoss())
loss = lossfunc(model(X), y)

with backpack(SumGradSquared()):
    loss.backward()

for param in model.parameters():
    print(param.grad)
    print(param.sum_grad_squared)

"""
Compute the gradient with Pytorch
and the diagonal of the Gauss-Newton with BackPACK
"""
from torch.nn import CrossEntropyLoss, Linear
from backpack.utils.examples import load_one_batch_mnist
from backpack import extend, backpack
from backpack.extensions import DiagGGNExact

X, y = load_one_batch_mnist(flat=True)
model = extend(Linear(784, 10))
lossfunc = extend(CrossEntropyLoss())
loss = lossfunc(model(X), y)

with backpack(DiagGGNExact()):
    loss.backward()

for param in model.parameters():
    print(param.grad)
    print(param.diag_ggn_exact)

"""
Compute the gradient with Pytorch
and KFAC with BackPACK
"""
from torch.nn import CrossEntropyLoss, Linear
from backpack.utils.examples import load_one_batch_mnist
from backpack import extend, backpack
from backpack.extensions import KFAC

X, y = load_one_batch_mnist(flat=True)
model = extend(Linear(784, 10))
lossfunc = extend(CrossEntropyLoss())
loss = lossfunc(model(X), y)

with backpack(KFAC()):
    loss.backward()

for param in model.parameters():
    print(param.grad)
    print(param.kfac)