BackPACK

"""
Compute the gradient with Pytorch

"""
from torch.nn import CrossEntropyLoss, Linear
from utils import load_mnist_data


X, y = load_mnist_data()
model = Linear(764, 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 utils import load_mnist_data
from backpack import extend, backpack, Variance

X, y = load_mnist_data()
model = extend(Linear(764, 10))
lossfunc = extend(CrossEntropyLoss())
loss = lossfunc(model(X), y)

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

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

"""
Compute the gradient with Pytorch
and the second moment with BackPACK
"""
from torch.nn import CrossEntropyLoss, Linear
from utils import load_mnist_data
from backpack import extend, backpack, SecondMoment

X, y = load_mnist_data()
model = extend(Linear(764, 10))
lossfunc = extend(CrossEntropyLoss())
loss = lossfunc(model(X), y)

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

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

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

X, y = load_mnist_data()
model = extend(Linear(764, 10))
lossfunc = extend(CrossEntropyLoss())
loss = lossfunc(model(X), y)

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

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

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

X, y = load_mnist_data()
model = extend(Linear(764, 10))
lossfunc = extend(CrossEntropyLoss())
loss = lossfunc(model(X), y)

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

for param in model.parameters():
    print(param.grad)
    print(param.kfac1, param.kfac2)