def posterior(
self, X: Tensor, observation_noise: bool = False, **kwargs: Any
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `(batch_shape) x q x d`-dim Tensor, where `d` is the dimension of the
feature space and `q` is the number of points considered jointly.
observation_noise: If True, add observation noise to the posterior.
detach_test_caches: If True, detach GPyTorch test caches during
computation of the posterior. Required for being able to compute
derivatives with respect to training inputs at test time (used
e.g. by qNoisyExpectedImprovement). Defaults to `True`.
Returns:
A `GPyTorchPosterior` object, representing a batch of `b` joint
distributions over `q` points. Includes observation noise if
`observation_noise=True`.
"""
self.eval() # make sure model is in eval mode
detach_test_caches = kwargs.get("detach_test_caches", True)
with ExitStack() as es:
es.enter_context(settings.debug(False))
es.enter_context(settings.fast_pred_var())
es.enter_context(settings.detach_test_caches(detach_test_caches))
mvn = self(X)
if observation_noise:
# TODO: Allow passing in observation noise via kwarg
mvn = self.likelihood(mvn, X)
return GPyTorchPosterior(mvn=mvn)
def posterior(self,
X: Tensor,
observation_noise: bool = False,
**kwargs: Any) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `(batch_shape) x q x d`-dim Tensor, where `d` is the dimension of the
feature space and `q` is the number of points considered jointly.
observation_noise: If True, add observation noise to the posterior.
Returns:
A `GPyTorchPosterior` object, representing a batch of `b` joint
distributions over `q` points. Includes observation noise if
`observation_noise=True`.
"""
self.eval() # make sure model is in eval mode
with ExitStack() as es:
es.enter_context(gpt_settings.debug(False))
es.enter_context(gpt_settings.fast_pred_var())
es.enter_context(
gpt_settings.detach_test_caches(
settings.propagate_grads.off()))
mvn = self(X)
if observation_noise:
# TODO: Allow passing in observation noise via kwarg
mvn = self.likelihood(mvn, X)
return GPyTorchPosterior(mvn=mvn)
def posterior(
self,
X: Tensor,
output_indices: Optional[List[int]] = None,
observation_noise: bool = False,
**kwargs: Any,
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `q x d` or `batch_shape x q x d` (batch mode) tensor, where `d` is the
dimension of the feature space (not including task indices) and
`q` is the number of points considered jointly.
output_indices: A list of indices, corresponding to the outputs over
which to compute the posterior (if the model is multi-output).
Can be used to speed up computation if only a subset of the
model's outputs are required for optimization. If omitted,
computes the posterior over all model outputs.
observation_noise: If True, add observation noise to the posterior.
detach_test_caches: If True, detach GPyTorch test caches during
computation of the posterior. Required for being able to compute
derivatives with respect to training inputs at test time (used
e.g. by qNoisyExpectedImprovement).
Returns:
A `GPyTorchPosterior` object, representing `batch_shape` joint
distributions over `q` points and the outputs selected by
`output_indices`. Includes measurement noise if
`observation_noise=True`.
"""
if output_indices is None:
output_indices = self._output_tasks
if any(i not in self._output_tasks for i in output_indices):
raise ValueError("Too many output indices")
# construct evaluation X
X_full = _make_X_full(X=X,
output_indices=output_indices,
tf=self._task_feature)
self.eval() # make sure model is in eval mode
detach_test_caches = kwargs.get("detach_test_caches", True)
with ExitStack() as es:
es.enter_context(settings.debug(False))
es.enter_context(settings.fast_pred_var())
es.enter_context(settings.detach_test_caches(detach_test_caches))
mvn = self(X_full)
if observation_noise:
# TODO: Allow passing in observation noise via kwarg
mvn = self.likelihood(mvn, X_full)
# If single-output, return the posterior of a single-output model
if len(output_indices) == 1:
return GPyTorchPosterior(mvn=mvn)
# Otherwise, make a MultitaskMultivariateNormal out of this
mtmvn = MultitaskMultivariateNormal(
mean=mvn.mean.view(*X.shape[:-1], len(output_indices)),
covariance_matrix=mvn.lazy_covariance_matrix,
interleaved=False,
)
return GPyTorchPosterior(mvn=mtmvn)
def posterior(
self, X: Tensor, observation_noise: bool = False, **kwargs: Any
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `(batch_shape) x q x d`-dim Tensor, where `d` is the dimension of the
feature space and `q` is the number of points considered jointly.
observation_noise: If True, add observation noise to the posterior.
detach_test_caches: If True, detach GPyTorch test caches during
computation of the posterior. Required for being able to compute
derivatives with respect to training inputs at test time (used
e.g. by qNoisyExpectedImprovement). Defaults to `True`.
Returns:
A `GPyTorchPosterior` object, representing a batch of `b` joint
distributions over `q` points. Includes observation noise if
`observation_noise=True`.
"""
self.eval() # make sure model is in eval mode
detach_test_caches = kwargs.get("detach_test_caches", True)
with ExitStack() as es:
es.enter_context(settings.debug(False))
es.enter_context(settings.fast_pred_var())
es.enter_context(settings.detach_test_caches(detach_test_caches))
mvn = self(X)
if observation_noise:
# TODO: Allow passing in observation noise via kwarg
mvn = self.likelihood(mvn, X)
return GPyTorchPosterior(mvn=mvn)
def gpt_posterior_settings():
r"""Context manager for settings used for computing model posteriors."""
with ExitStack() as es:
es.enter_context(gpt_settings.debug(False))
es.enter_context(gpt_settings.fast_pred_var())
es.enter_context(
gpt_settings.detach_test_caches(settings.propagate_grads.off()))
yield
def posterior(
self,
X: Tensor,
output_indices: Optional[List[int]] = None,
observation_noise: bool = False,
**kwargs: Any,
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `q x d` or `batch_shape x q x d` (batch mode) tensor, where `d` is the
dimension of the feature space (not including task indices) and
`q` is the number of points considered jointly.
output_indices: A list of indices, corresponding to the outputs over
which to compute the posterior (if the model is multi-output).
Can be used to speed up computation if only a subset of the
model's outputs are required for optimization. If omitted,
computes the posterior over all model outputs.
observation_noise: If True, add observation noise to the posterior.
detach_test_caches: If True, detach GPyTorch test caches during
computation of the posterior. Required for being able to compute
derivatives with respect to training inputs at test time (used
e.g. by qNoisyExpectedImprovement).
Returns:
A `GPyTorchPosterior` object, representing `batch_shape` joint
distributions over `q` points and the outputs selected by
`output_indices`. Includes measurement noise if
`observation_noise=True`.
"""
if output_indices is None:
output_indices = self._output_tasks
if any(i not in self._output_tasks for i in output_indices):
raise ValueError("Too many output indices")
# construct evaluation X
X_full = _make_X_full(X=X, output_indices=output_indices, tf=self._task_feature)
self.eval() # make sure model is in eval mode
detach_test_caches = kwargs.get("detach_test_caches", True)
with ExitStack() as es:
es.enter_context(settings.debug(False))
es.enter_context(settings.fast_pred_var())
es.enter_context(settings.detach_test_caches(detach_test_caches))
mvn = self(X_full)
if observation_noise:
# TODO: Allow passing in observation noise via kwarg
mvn = self.likelihood(mvn, X_full)
# If single-output, return the posterior of a single-output model
if len(output_indices) == 1:
return GPyTorchPosterior(mvn=mvn)
# Otherwise, make a MultitaskMultivariateNormal out of this
mtmvn = MultitaskMultivariateNormal(
mean=mvn.mean.view(*X.shape[:-1], len(output_indices)),
covariance_matrix=mvn.lazy_covariance_matrix,
interleaved=False,
)
return GPyTorchPosterior(mvn=mtmvn)
def posterior(
self,
X: Tensor,
output_indices: Optional[List[int]] = None,
observation_noise: bool = False,
**kwargs: Any,
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `(batch_shape) x q x d`-dim Tensor, where `d` is the dimension of the
feature space and `q` is the number of points considered jointly.
output_indices: A list of indices, corresponding to the outputs over
which to compute the posterior (if the model is multi-output).
Can be used to speed up computation if only a subset of the
model's outputs are required for optimization. If omitted,
computes the posterior over all model outputs.
observation_noise: If True, add observation noise to the posterior.
Returns:
A `GPyTorchPosterior` object, representing `batch_shape` joint
distributions over `q` points and the outputs selected by
`output_indices` each. Includes observation noise if
`observation_noise=True`.
"""
self.eval() # make sure model is in eval mode
with ExitStack() as es:
es.enter_context(gpt_settings.debug(False))
es.enter_context(gpt_settings.fast_pred_var())
es.enter_context(
gpt_settings.detach_test_caches(
settings.propagate_grads.off()))
# insert a dimension for the output dimension
if self._num_outputs > 1:
X, output_dim_idx = add_output_dim(
X=X, original_batch_shape=self._input_batch_shape)
mvn = self(X)
if observation_noise:
if isinstance(self.likelihood, FixedNoiseGaussianLikelihood):
# Use the mean of the previous noise values (TODO: be smarter here).
noise = self.likelihood.noise.mean().expand(X.shape[:-1])
mvn = self.likelihood(mvn, X, noise=noise)
else:
mvn = self.likelihood(mvn, X)
if self._num_outputs > 1:
mean_x = mvn.mean
covar_x = mvn.covariance_matrix
output_indices = output_indices or range(self._num_outputs)
mvns = [
MultivariateNormal(
mean_x.select(dim=output_dim_idx, index=t),
lazify(covar_x.select(dim=output_dim_idx, index=t)),
) for t in output_indices
]
mvn = MultitaskMultivariateNormal.from_independent_mvns(
mvns=mvns)
return GPyTorchPosterior(mvn=mvn)
def posterior(
self,
X: Tensor,
output_indices: Optional[List[int]] = None,
observation_noise: bool = False,
**kwargs: Any,
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `b x q x d`-dim Tensor, where `d` is the dimension of the
feature space, `q` is the number of points considered jointly,
and `b` is the batch dimension.
output_indices: A list of indices, corresponding to the outputs over
which to compute the posterior (if the model is multi-output).
Can be used to speed up computation if only a subset of the
model's outputs are required for optimization. If omitted,
computes the posterior over all model outputs.
observation_noise: If True, add observation noise to the posterior.
Returns:
A `GPyTorchPosterior` object, representing `batch_shape` joint
distributions over `q` points and the outputs selected by
`output_indices` each. Includes measurement noise if
`observation_noise=True`.
"""
self.eval() # make sure model is in eval mode
with ExitStack() as es:
es.enter_context(gpt_settings.debug(False))
es.enter_context(gpt_settings.fast_pred_var())
es.enter_context(
gpt_settings.detach_test_caches(
settings.propagate_grads.off()))
if output_indices is not None:
mvns = [self.forward_i(i, X) for i in output_indices]
if observation_noise:
lh_kwargs = [{
"noise": lh.noise.mean().expand(X.shape[:-1])
} if isinstance(lh, FixedNoiseGaussianLikelihood) else {}
for lh in self.likelihood.likelihoods]
mvns = [
self.likelihood_i(i, mvn, X,
**lkws) for i, mvn, lkws in zip(
output_indices, mvns, lh_kwargs)
]
else:
mvns = self(*[X for _ in range(self.num_outputs)])
if observation_noise:
# TODO: Allow passing in observation noise via kwarg
mvns = self.likelihood(*[(mvn, X) for mvn in mvns])
if len(mvns) == 1:
return GPyTorchPosterior(mvn=mvns[0])
else:
return GPyTorchPosterior(
mvn=MultitaskMultivariateNormal.from_independent_mvns(
mvns=mvns))
def posterior(
self,
X: Tensor,
output_indices: Optional[List[int]] = None,
observation_noise: bool = False,
**kwargs: Any,
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `(batch_shape) x q x d`-dim Tensor, where `d` is the dimension of the
feature space and `q` is the number of points considered jointly.
output_indices: A list of indices, corresponding to the outputs over
which to compute the posterior (if the model is multi-output).
Can be used to speed up computation if only a subset of the
model's outputs are required for optimization. If omitted,
computes the posterior over all model outputs.
observation_noise: If True, add observation noise to the posterior.
propagate_grads: If True, do not detach GPyTorch's test caches when
computing of the posterior. Required for being able to compute
derivatives with respect to training inputs at test time (used
e.g. by qNoisyExpectedImprovement). Defaults to `False`.
Returns:
A `GPyTorchPosterior` object, representing `batch_shape` joint
distributions over `q` points and the outputs selected by
`output_indices` each. Includes observation noise if
`observation_noise=True`.
"""
self.eval() # make sure model is in eval mode
detach_test_caches = not kwargs.get("propagate_grads", False)
with ExitStack() as es:
es.enter_context(settings.debug(False))
es.enter_context(settings.fast_pred_var())
es.enter_context(settings.detach_test_caches(detach_test_caches))
# insert a dimension for the output dimension
if self._num_outputs > 1:
X, output_dim_idx = add_output_dim(
X=X, original_batch_shape=self._input_batch_shape
)
mvn = self(X)
if observation_noise:
mvn = self.likelihood(mvn, X)
if self._num_outputs > 1:
mean_x = mvn.mean
covar_x = mvn.covariance_matrix
output_indices = output_indices or range(self._num_outputs)
mvns = [
MultivariateNormal(
mean_x.select(dim=output_dim_idx, index=t),
lazify(covar_x.select(dim=output_dim_idx, index=t)),
)
for t in output_indices
]
mvn = MultitaskMultivariateNormal.from_independent_mvns(mvns=mvns)
return GPyTorchPosterior(mvn=mvn)
def posterior(
self,
X: Tensor,
output_indices: Optional[List[int]] = None,
observation_noise: bool = False,
**kwargs: Any,
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `b x q x d`-dim Tensor, where `d` is the dimension of the
feature space, `q` is the number of points considered jointly,
and `b` is the batch dimension.
output_indices: A list of indices, corresponding to the outputs over
which to compute the posterior (if the model is multi-output).
Can be used to speed up computation if only a subset of the
model's outputs are required for optimization. If omitted,
computes the posterior over all model outputs.
observation_noise: If True, add observation noise to the posterior.
detach_test_caches: If True, detach GPyTorch test caches during
computation of the posterior. Required for being able to compute
derivatives with respect to training inputs at test time (used
e.g. by qNoisyExpectedImprovement).
Returns:
A `GPyTorchPosterior` object, representing `batch_shape` joint
distributions over `q` points and the outputs selected by
`output_indices` each. Includes measurement noise if
`observation_noise=True`.
"""
detach_test_caches = kwargs.get("detach_test_caches", True)
self.eval() # make sure model is in eval mode
with ExitStack() as es:
es.enter_context(settings.debug(False))
es.enter_context(settings.fast_pred_var())
es.enter_context(settings.detach_test_caches(detach_test_caches))
if output_indices is not None:
mvns = [self.forward_i(i, X) for i in output_indices]
if observation_noise:
mvns = [
self.likelihood_i(i, mvn, X)
for i, mvn in zip(output_indices, mvns)
]
else:
mvns = self(*[X for _ in range(self.num_outputs)])
if observation_noise:
# TODO: Allow passing in observation noise via kwarg
mvns = self.likelihood(*[(mvn, X) for mvn in mvns])
if len(mvns) == 1:
return GPyTorchPosterior(mvn=mvns[0])
else:
return GPyTorchPosterior(
mvn=MultitaskMultivariateNormal.from_independent_mvns(mvns=mvns)
)
def posterior(
self,
X: Tensor,
output_indices: Optional[List[int]] = None,
observation_noise: bool = False,
**kwargs: Any,
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `(batch_shape) x q x d`-dim Tensor, where `d` is the dimension of the
feature space and `q` is the number of points considered jointly.
output_indices: A list of indices, corresponding to the outputs over
which to compute the posterior (if the model is multi-output).
Can be used to speed up computation if only a subset of the
model's outputs are required for optimization. If omitted,
computes the posterior over all model outputs.
observation_noise: If True, add observation noise to the posterior.
detach_test_caches: If True, detach GPyTorch test caches during
computation of the posterior. Required for being able to compute
derivatives with respect to training inputs at test time (used
e.g. by qNoisyExpectedImprovement). Defaults to `True`.
Returns:
A `GPyTorchPosterior` object, representing `batch_shape` joint
distributions over `q` points and the outputs selected by
`output_indices` each. Includes observation noise if
`observation_noise=True`.
"""
self.eval() # make sure model is in eval mode
detach_test_caches = kwargs.get("detach_test_caches", True)
with ExitStack() as es:
es.enter_context(settings.debug(False))
es.enter_context(settings.fast_pred_var())
es.enter_context(settings.detach_test_caches(detach_test_caches))
# insert a dimension for the output dimension
if self._num_outputs > 1:
X, output_dim_idx = add_output_dim(
X=X, original_batch_shape=self._input_batch_shape
)
mvn = self(X)
mean_x = mvn.mean
covar_x = mvn.covariance_matrix
if self._num_outputs > 1:
output_indices = output_indices or range(self._num_outputs)
mvns = [
MultivariateNormal(
mean_x.select(dim=output_dim_idx, index=t),
lazify(covar_x.select(dim=output_dim_idx, index=t)),
)
for t in output_indices
]
mvn = MultitaskMultivariateNormal.from_independent_mvns(mvns=mvns)
return GPyTorchPosterior(mvn=mvn)
def posterior(
self,
X: Tensor,
output_indices: Optional[List[int]] = None,
observation_noise: bool = False,
**kwargs: Any,
) -> GPyTorchPosterior:
r"""Computes the posterior over model outputs at the provided points.
Args:
X: A `b x q x d`-dim Tensor, where `d` is the dimension of the
feature space, `q` is the number of points considered jointly,
and `b` is the batch dimension.
output_indices: A list of indices, corresponding to the outputs over
which to compute the posterior (if the model is multi-output).
Can be used to speed up computation if only a subset of the
model's outputs are required for optimization. If omitted,
computes the posterior over all model outputs.
observation_noise: If True, add observation noise to the posterior.
detach_test_caches: If True, detach GPyTorch test caches during
computation of the posterior. Required for being able to compute
derivatives with respect to training inputs at test time (used
e.g. by qNoisyExpectedImprovement).
Returns:
A `GPyTorchPosterior` object, representing `batch_shape` joint
distributions over `q` points and the outputs selected by
`output_indices` each. Includes measurement noise if
`observation_noise=True`.
"""
detach_test_caches = kwargs.get("detach_test_caches", True)
self.eval() # make sure model is in eval mode
with ExitStack() as es:
es.enter_context(settings.debug(False))
es.enter_context(settings.fast_pred_var())
es.enter_context(settings.detach_test_caches(detach_test_caches))
if output_indices is not None:
mvns = [self.forward_i(i, X) for i in output_indices]
if observation_noise:
mvns = [
self.likelihood_i(i, mvn, X)
for i, mvn in zip(output_indices, mvns)
]
else:
mvns = self(*[X for _ in range(self.num_outputs)])
if observation_noise:
# TODO: Allow passing in observation noise via kwarg
mvns = self.likelihood(*[(mvn, X) for mvn in mvns])
if len(mvns) == 1:
return GPyTorchPosterior(mvn=mvns[0])
else:
return GPyTorchPosterior(
mvn=MultitaskMultivariateNormal.from_independent_mvns(mvns=mvns)
)
acq_value.item(),
pred_rmse.item(),
pred_avg_variance.item()
]
print("Step RMSE: ", pred_rmse)
all_outputs.append(step_output_list)
start_ind = end_ind
end_ind = int(end_ind + args.batch_size)
output_dict = {
"model_state_dict": model.cpu().state_dict(),
"queried_points": {
'x': model.cpu().train_inputs[0],
'y': model.cpu().train_targets
},
"results": DataFrame(all_outputs)
}
torch.save(output_dict, args.output)
if __name__ == "__main__":
args = parse()
with fast_pred_var(True), \
use_toeplitz(args.toeplitz), \
detach_test_caches(True), \
max_cholesky_size(args.cholesky_size), \
max_root_decomposition_size(args.sketch_size), \
root_pred_var(True):
main(args)
def main(args):
if args.cuda and torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
init_dict, train_dict, test_dict = prepare_data(args.data_loc,
args.num_init,
args.num_total,
test_is_year=False)
init_x, init_y, init_y_var = (
init_dict["x"].to(device),
init_dict["y"].to(device),
init_dict["y_var"].to(device),
)
train_x, train_y, train_y_var = (
train_dict["x"].to(device),
train_dict["y"].to(device),
train_dict["y_var"].to(device),
)
test_x, test_y, test_y_var = (
test_dict["x"].to(device),
test_dict["y"].to(device),
test_dict["y_var"].to(device),
)
model = FixedNoiseOnlineSKIGP(
init_x,
init_y.view(-1, 1),
init_y_var.view(-1, 1),
GridInterpolationKernel(
base_kernel=ScaleKernel(
MaternKernel(
ard_num_dims=2,
nu=0.5,
lengthscale_prior=GammaPrior(3.0, 6.0),
),
outputscale_prior=GammaPrior(2.0, 0.15),
),
grid_size=30,
num_dims=2,
grid_bounds=torch.tensor([[0.0, 1.0], [0.0, 1.0]]),
),
learn_additional_noise=False,
).to(device)
mll = BatchedWoodburyMarginalLogLikelihood(model.likelihood, model)
print("---- Fitting initial model ----")
start = time.time()
with skip_logdet_forward(True), max_root_decomposition_size(
args.sketch_size), use_toeplitz(args.toeplitz):
fit_gpytorch_torch(mll, options={"lr": 0.1, "maxiter": 1000})
end = time.time()
print("Elapsed fitting time: ", end - start)
model.zero_grad()
model.eval()
print("--- Generating initial predictions on test set ----")
start = time.time()
with detach_test_caches(True), max_root_decomposition_size(
args.sketch_size), max_cholesky_size(
args.cholesky_size), use_toeplitz(args.toeplitz):
pred_dist = model(test_x)
pred_mean = pred_dist.mean.detach()
# pred_var = pred_dist.variance.detach()
end = time.time()
print("Elapsed initial prediction time: ", end - start)
rmse_initial = ((pred_mean.view(-1) - test_y.view(-1))**2).mean().sqrt()
print("Initial RMSE: ", rmse_initial.item())
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
mll_time_list = []
rmse_list = []
for i in range(500, train_x.shape[0]):
model.zero_grad()
model.train()
start = time.time()
with skip_logdet_forward(True), max_root_decomposition_size(
args.sketch_size), max_cholesky_size(
args.cholesky_size), use_toeplitz(args.toeplitz):
loss = -mll(model(train_x[:i]), train_y[:i]).sum()
loss.backward()
mll_time = start - time.time()
optimizer.step()
model.zero_grad()
optimizer.zero_grad()
start = time.time()
with torch.no_grad():
model.condition_on_observations(
train_x[i].unsqueeze(0),
train_y[i].view(1, 1),
train_y_var[i].view(-1, 1),
inplace=True,
)
fantasy_time = start - time.time()
mll_time_list.append([-mll_time, -fantasy_time])
if i % 25 == 0:
start = time.time()
model.eval()
model.zero_grad()
with detach_test_caches(), max_root_decomposition_size(
args.sketch_size), max_cholesky_size(args.cholesky_size):
pred_dist = model(test_x)
end = time.time()
rmse = (((pred_dist.mean -
test_y.view(-1))**2).mean().sqrt().item())
rmse_list.append([rmse, end - start])
print("Current RMSE: ", rmse)
print("Outputscale: ",
model.covar_module.base_kernel.raw_outputscale)
print(
"Lengthscale: ",
model.covar_module.base_kernel.base_kernel.raw_lengthscale,
)
print("Step: ", i, "Train Loss: ", loss)
optimizer.param_groups[0]["lr"] *= 0.9
torch.save({
"training": mll_time_list,
"predictions": rmse_list
}, args.output)
