def test_interface_variance(): interface_test(new_ext.Variance())
model = Sequential(Flatten(), Linear(784, 10),) lossfunc = CrossEntropyLoss() model = extend(model) lossfunc = extend(lossfunc) # %% # We can now evaluate the loss and do a backward pass with Backpack # ----------------------------------------------------------------- loss = lossfunc(model(X), y) with backpack( extensions.BatchGrad(), extensions.Variance(), extensions.SumGradSquared(), extensions.BatchL2Grad(), extensions.DiagGGNMC(mc_samples=1), extensions.DiagGGNExact(), extensions.DiagHessian(), extensions.KFAC(mc_samples=1), extensions.KFLR(), extensions.KFRA(), ): loss.backward() # %% # And here are the results # -----------------------------------------------------------------
def variance(self): with backpack(new_ext.Variance()): _, _, loss = self.problem.forward_pass() loss.backward() variances = [p.variance for p in self.problem.model.parameters()] return variances
def variance(self): with backpack(new_ext.Variance()): self.loss().backward() variances = [p.variance for p in self.model.parameters()] return variances
""" from torch.nn import CrossEntropyLoss, Flatten, Linear, Sequential from backpack import backpack, extend, extensions from backpack.utils.examples import load_mnist_data B = 4 X, y = load_mnist_data(B) print("# Gradient with PyTorch, gradient variance with BackPACK | B =", B) model = Sequential( Flatten(), Linear(784, 10), ) lossfunc = CrossEntropyLoss() model = extend(model) lossfunc = extend(lossfunc) loss = lossfunc(model(X), y) with backpack(extensions.Variance()): loss.backward() for name, param in model.named_parameters(): print(name) print(".grad.shape: ", param.grad.shape) print(".variance.shape: ", param.variance.shape)
def variance(self) -> List[Tensor]: # noqa:D102 with backpack(new_ext.Variance()): _, _, loss = self.problem.forward_pass() loss.backward() return self.problem.collect_data("variance")