def set_params_numpy( self, model: Any, new_param_array: torch.Tensor, ) -> None: """ Sets the model's parameters using the values in `new_param_array`. Parameters ---------- new_param_array : 1D ndarray. Should have one element for each element of the tensors in `self.parameters`. Returns ------- None. """ # Get the property dictionary for this module _, property_dict, _ = self.get_params_numpy(model) # Set the parameters numpy_converter.set_params_with_array(model, new_param_array, property_dict)
def test_set_parameters(self): for dtype in (torch.float, torch.double): # get a test module train_x = torch.tensor([[1.0, 2.0, 3.0]], device=self.device, dtype=dtype) train_y = torch.tensor([4.0], device=self.device, dtype=dtype) likelihood = GaussianLikelihood() model = ExactGP(train_x, train_y, likelihood) model.covar_module = RBFKernel(ard_num_dims=3) model.mean_module = ConstantMean() model.to(device=self.device, dtype=dtype) mll = ExactMarginalLogLikelihood(likelihood, model) # get parameters x, pdict, bounds = module_to_array(module=mll) # Set parameters mll = set_params_with_array(mll, np.array([1.0, 2.0, 3.0, 4.0, 5.0]), pdict) z = dict(mll.named_parameters()) self.assertTrue( torch.equal( z["likelihood.noise_covar.raw_noise"], torch.tensor([1.0], device=self.device, dtype=dtype), )) self.assertTrue( torch.equal( z["model.covar_module.raw_lengthscale"], torch.tensor([[2.0, 3.0, 4.0]], device=self.device, dtype=dtype), )) self.assertTrue( torch.equal( z["model.mean_module.constant"], torch.tensor([5.0], device=self.device, dtype=dtype), )) # Extract again x2, pdict2, bounds2 = module_to_array(module=mll) self.assertTrue( np.array_equal(x2, np.array([1.0, 2.0, 3.0, 4.0, 5.0])))
def _scipy_objective_and_grad( x: np.ndarray, mll: MarginalLogLikelihood, property_dict: Dict[str, TorchAttr]) -> Tuple[float, np.ndarray]: r"""Get objective and gradient in format that scipy expects. Args: x: The (flattened) input parameters. mll: The MarginalLogLikelihood module to evaluate. property_dict: The property dictionary required to "unflatten" the input parameter vector, as generated by `module_to_array`. Returns: 2-element tuple containing - The objective value. - The gradient of the objective. """ mll = set_params_with_array(mll, x, property_dict) train_inputs, train_targets = mll.model.train_inputs, mll.model.train_targets mll.zero_grad() try: # catch linear algebra errors in gpytorch output = mll.model(*train_inputs) args = [output, train_targets] + _get_extra_mll_args(mll) loss = -mll(*args).sum() except RuntimeError as e: if isinstance(e, NotPSDError): raise e if isinstance(e, NanError) or "singular" in e.args[0]: return float("nan"), np.full_like(x, "nan") raise e # pragma: nocover loss.backward() param_dict = OrderedDict(mll.named_parameters()) grad = [] for p_name in property_dict: t = param_dict[p_name].grad if t is None: # this deals with parameters that do not affect the loss grad.append(np.zeros(property_dict[p_name].shape.numel())) else: grad.append(t.detach().view(-1).cpu().double().clone().numpy()) mll.zero_grad() return loss.item(), np.concatenate(grad)
def test_set_parameters(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") for dtype in (torch.float, torch.double): # get a test module train_x = torch.tensor([[1.0, 2.0, 3.0]], device=device, dtype=dtype) train_y = torch.tensor([4.0], device=device, dtype=dtype) likelihood = GaussianLikelihood() model = ExactGP(train_x, train_y, likelihood) model.covar_module = RBFKernel(ard_num_dims=3) model.mean_module = ConstantMean() model.to(device=device, dtype=dtype) mll = ExactMarginalLogLikelihood(likelihood, model) # get parameters x, pdict, bounds = module_to_array(module=mll) # Set parameters mll = set_params_with_array(mll, np.array([1.0, 2.0, 3.0, 4.0, 5.0]), pdict) z = dict(mll.named_parameters()) self.assertTrue( torch.equal( z["likelihood.noise_covar.raw_noise"], torch.tensor([1.0], device=device, dtype=dtype), ) ) self.assertTrue( torch.equal( z["model.covar_module.raw_lengthscale"], torch.tensor([[2.0, 3.0, 4.0]], device=device, dtype=dtype), ) ) self.assertTrue( torch.equal( z["model.mean_module.constant"], torch.tensor([5.0], device=device, dtype=dtype), ) ) # Extract again x2, pdict2, bounds2 = module_to_array(module=mll) self.assertTrue(np.array_equal(x2, np.array([1.0, 2.0, 3.0, 4.0, 5.0])))
def fit_gpytorch_scipy( mll: MarginalLogLikelihood, bounds: Optional[ParameterBounds] = None, method: str = "L-BFGS-B", options: Optional[Dict[str, Any]] = None, track_iterations: bool = True, ) -> Tuple[MarginalLogLikelihood, Dict[str, Union[ float, List[OptimizationIteration]]]]: r"""Fit a gpytorch model by maximizing MLL with a scipy optimizer. The model and likelihood in mll must already be in train mode. Note: this method requires that the model has `train_inputs` and `train_targets`. Args: mll: MarginalLogLikelihood to be maximized. bounds: A dictionary mapping parameter names to tuples of lower and upper bounds. method: Solver type, passed along to scipy.minimize. options: Dictionary of solver options, passed along to scipy.minimize. track_iterations: Track the function values and wall time for each iteration. Returns: 2-element tuple containing - MarginalLogLikelihood with parameters optimized in-place. - Dictionary with the following key/values: "fopt": Best mll value. "wall_time": Wall time of fitting. "iterations": List of OptimizationIteration objects with information on each iteration. If track_iterations is False, will be empty. Example: >>> gp = SingleTaskGP(train_X, train_Y) >>> mll = ExactMarginalLogLikelihood(gp.likelihood, gp) >>> mll.train() >>> fit_gpytorch_scipy(mll) >>> mll.eval() """ options = options or {} x0, property_dict, bounds = module_to_array(module=mll, bounds=bounds, exclude=options.pop( "exclude", None)) x0 = x0.astype(np.float64) if bounds is not None: bounds = Bounds(lb=bounds[0], ub=bounds[1], keep_feasible=True) xs = [] ts = [] t1 = time.time() def store_iteration(xk): xs.append(xk.copy()) ts.append(time.time() - t1) cb = store_iteration if track_iterations else None res = minimize( _scipy_objective_and_grad, x0, args=(mll, property_dict), bounds=bounds, method=method, jac=True, options=options, callback=cb, ) iterations = [] if track_iterations: for i, xk in enumerate(xs): obj, _ = _scipy_objective_and_grad(x=xk, mll=mll, property_dict=property_dict) iterations.append(OptimizationIteration(i, obj, ts[i])) # Construct info dict info_dict = { "fopt": float(res.fun), "wall_time": time.time() - t1, "iterations": iterations, } if not res.success: try: # Some res.message are bytes msg = res.message.decode("ascii") except AttributeError: # Others are str msg = res.message warnings.warn(f"Fitting failed with the optimizer reporting '{msg}'", OptimizationWarning) # Set to optimum mll = set_params_with_array(mll, res.x, property_dict) return mll, info_dict