def hvp_fn(x): x_tensor = torch.tensor(x, requires_grad=False) if gpu >= 0: x_tensor = x_tensor.cuda() params, names = make_functional(model) # Make params regular Tensors instead of nn.Parameter params = tuple(p.detach().requires_grad_() for p in params) flat_params = parameters_to_vector(params) hvp = torch.zeros_like(flat_params) for x_train, y_train in train_loader: if gpu >= 0: x_train, y_train = x_train.cuda(), y_train.cuda() def f(flat_params_): split_params = tensor_to_tuple(flat_params_, params) load_weights(model, names, split_params) out = model(x_train) loss = calc_loss(out, y_train) return loss batch_hvp = vhp(f, flat_params, x_tensor, strict=True)[1] hvp += batch_hvp / float(len(train_loader)) with torch.no_grad(): load_weights(model, names, params, as_params=True) damped_hvp = hvp + damp * v_flat return damped_hvp.cpu().numpy()
def s_test(x_test, y_test, model, i, samples_loader, gpu=-1, damp=0.01, scale=25.0): """s_test can be precomputed for each test point of interest, and then multiplied with grad_z to get the desired value for each training point. Here, stochastic estimation is used to calculate s_test. s_test is the Inverse Hessian Vector Product. Arguments: x_test: torch tensor, test data points, such as test images y_test: torch tensor, contains all test data labels model: torch NN, model used to evaluate the dataset i: the sample number samples_loader: torch DataLoader, can load the training dataset gpu: int, GPU id to use if >=0 and -1 means use CPU damp: float, dampening factor scale: float, scaling factor Returns: h_estimate: list of torch tensors, s_test""" v = grad_z(x_test, y_test, model, gpu) h_estimate = v params, names = make_functional(model) # Make params regular Tensors instead of nn.Parameter params = tuple(p.detach().requires_grad_() for p in params) # TODO: Dynamically set the recursion depth so that iterations stop once h_estimate stabilises progress_bar = tqdm(samples_loader, desc=f"IHVP sample {i}") for i, (x_train, y_train) in enumerate(progress_bar): if gpu >= 0: x_train, y_train = x_train.cuda(), y_train.cuda() def f(*new_params): load_weights(model, names, new_params) out = model(x_train) loss = model.loss(out, y_train) #calc_loss(out, y_train, loss_func=loss_func) return loss hv = vhp(f, params, tuple(h_estimate), strict=True)[1] # Recursively calculate h_estimate with torch.no_grad(): h_estimate = [ _v + (1 - damp) * _h_e - _hv / scale for _v, _h_e, _hv in zip(v, h_estimate, hv) ] if i % 100 == 0: norm = sum([h_.norm() for h_ in h_estimate]) progress_bar.set_postfix({"est_norm": norm.item()}) with torch.no_grad(): load_weights(model, names, params, as_params=True) return h_estimate
def s_test_sample_exact(model, x_test, y_test, train_loader, gpu=-1): grads = grad_z(x_test, y_test, model, gpu=gpu) flat_grads = parameters_to_vector(grads) def make_loss_f(model, params, names, x, y): def f(flat_params_): split_params = tensor_to_tuple(flat_params_, params) load_weights(model, names, split_params) out = model(x) loss = model.loss(out, y) return loss return f # Make model functional params, names = make_functional(model) # Make params regular Tensors instead of nn.Parameter params = tuple(p.detach().requires_grad_() for p in params) flat_params = parameters_to_vector(params) h = torch.zeros([flat_params.shape[0], flat_params.shape[0]]) if gpu >= 0: h = h.cuda() # Compute real IHVP for x_train, y_train in train_loader: if gpu >= 0: x_train, y_train = x_train.cuda(), y_train.cuda() f = make_loss_f(model, params, names, x_train, y_train) batch_h = hessian(f, flat_params, strict=True) with torch.no_grad(): h += batch_h / float(len(train_loader)) h = (h + h.transpose(0,1))/2 with torch.no_grad(): load_weights(model, names, params, as_params=True) inv_h = torch.inverse(h) print("Inverse Hessian") print(inv_h) real_ihvp = inv_h @ flat_grads return tensor_to_tuple(real_ihvp, params)
def setUpClass(cls) -> None: pl.seed_everything(0) cls.n_features = 10 cls.n_params = 2 * cls.n_features cls.model = LinearRegression(cls.n_features) gpus = 1 if torch.cuda.is_available() else 0 trainer = pl.Trainer(gpus=gpus, max_epochs=10) # trainer.fit(cls.model) print(tuple(cls.model.parameters())) use_sklearn = True if use_sklearn: train_dataset = DummyDataset(cls.n_features) clf = SklearnLR() clf.fit(train_dataset.data, train_dataset.targets) with torch.no_grad(): cls.model.linear.weight = torch.nn.Parameter( torch.tensor([clf.coef_], dtype=torch.float)) cls.model.linear.bias = torch.nn.Parameter( torch.tensor([clf.intercept_], dtype=torch.float)) cls.train_loader = cls.model.train_dataloader(batch_size=40000) # Setup test point data cls.test_idx = 8 cls.x_test = torch.tensor([cls.model.test_set.data[[cls.test_idx]]], dtype=torch.float) cls.y_test = torch.tensor([cls.model.test_set.targets[[cls.test_idx]]], dtype=torch.float) # Compute estimated IVHP cls.gpu = 1 if torch.cuda.is_available() else -1 # Compute anc flatten grad grads = grad_z(cls.x_test, cls.y_test, cls.model, gpu=cls.gpu) flat_grads = parameters_to_vector(grads) print("Grads:") print(flat_grads) # Make model functional params, names = make_functional(cls.model) # Make params regular Tensors instead of nn.Parameter params = tuple(p.detach().requires_grad_() for p in params) flat_params = parameters_to_vector(params) # Initialize Hessian h = torch.zeros([flat_params.shape[0], flat_params.shape[0]]) # Compute real IHVP for x_train, y_train in cls.train_loader: if cls.gpu >= 0: x_train, y_train = x_train.cuda(), y_train.cuda() def f(flat_params_): split_params = tensor_to_tuple(flat_params_, params) load_weights(cls.model, names, split_params) out = cls.model(x_train) loss = calc_loss(out, y_train) return loss batch_h = hessian(f, flat_params, strict=True) with torch.no_grad(): h += batch_h / float(len(cls.train_loader)) print("Hessian:") print(h) complete_x_train = cls.train_loader.dataset.data real_hessian = complete_x_train.T @ complete_x_train / complete_x_train.shape[ 0] * 2 print(real_hessian) print(np.linalg.norm(real_hessian - h.cpu().numpy()[:10, :10])) np.save("hessian_pytorch.npy", h.cpu().numpy()) # Make the model back `nn` with torch.no_grad(): load_weights(cls.model, names, params, as_params=True) inv_h = torch.inverse(h) print("Inverse Hessian") print(inv_h) cls.real_ihvp = inv_h @ flat_grads print("Real IHVP") print(cls.real_ihvp)
def setUpClass(cls) -> None: pl.seed_everything(0) cls.n_features = 10 cls.n_classes = 3 cls.n_params = cls.n_classes * cls.n_features + cls.n_features cls.wd = wd = 1e-2 # weight decay=1/(nC) cls.model = LogisticRegression(cls.n_classes, cls.n_features, wd=cls.wd) gpus = 1 if torch.cuda.is_available() else 0 trainer = pl.Trainer(gpus=gpus, max_epochs=10) # trainer.fit(self.model) use_sklearn = True if use_sklearn: cls.train_dataset = cls.model.training_set #DummyDataset(cls.n_features, cls.n_classes) multi_class = "multinomial" if cls.model.n_classes != 2 else "auto" clf = SklearnLogReg(C=1 / len(cls.train_dataset) / wd, tol=1e-8, max_iter=1000, multi_class=multi_class) clf.fit(cls.train_dataset.data, cls.train_dataset.targets) with torch.no_grad(): cls.model.linear.weight = torch.nn.Parameter( torch.tensor(clf.coef_, dtype=torch.float)) cls.model.linear.bias = torch.nn.Parameter( torch.tensor(clf.intercept_, dtype=torch.float)) # Setup test point data cls.test_idx = 5 cls.x_test = torch.tensor(cls.model.test_set.data[[cls.test_idx]], dtype=torch.float) cls.y_test = torch.tensor(cls.model.test_set.targets[[cls.test_idx]], dtype=torch.long) # Compute estimated IVHP cls.gpu = 1 if torch.cuda.is_available() else -1 if cls.gpu >= 0: cls.model = cls.model.cuda() cls.x_test = cls.x_test.cuda() cls.y_test = cls.y_test.cuda() cls.train_loader = cls.model.train_dataloader(batch_size=40000) # Compute anc flatten grad grads = grad_z(cls.x_test, cls.y_test, cls.model, gpu=cls.gpu) flat_grads = parameters_to_vector(grads) print("Grads:") print(flat_grads) # Make model functional params, names = make_functional(cls.model) # Make params regular Tensors instead of nn.Parameter params = tuple(p.detach().requires_grad_() for p in params) flat_params = parameters_to_vector(params) # Initialize Hessian h = torch.zeros([flat_params.shape[0], flat_params.shape[0]]) if cls.gpu == 1: h = h.cuda() # Compute real IHVP for x_train, y_train in cls.train_loader: if cls.gpu >= 0: x_train, y_train = x_train.cuda(), y_train.cuda() f = make_loss_f(cls.model, params, names, x_train, y_train, wd=wd) batch_h = hessian(f, flat_params, strict=True) with torch.no_grad(): h += batch_h / float(len(cls.train_loader)) h = (h + h.transpose(0, 1)) / 2 print("Hessian:") print(h) np.save("hessian_pytorch.npy", h.cpu().numpy()) from numpy import linalg as LA ei = LA.eig(h.cpu().numpy())[0] print('ei=', ei) print("max,min eigen value=", ei.max(), ei.min()) assert ei.min() > 0, "Error: Non-positive Eigenvalues" # Make the model back `nn` with torch.no_grad(): load_weights(cls.model, names, params, as_params=True) inv_h = torch.inverse(h) print("Inverse Hessian") print(inv_h) cls.real_ihvp = inv_h @ flat_grads print("Real IHVP") print(cls.real_ihvp)