Continue Warmup / Sampling after Interruption

General
1

Hello,

I want to run an IRT model on a relatively large dataset (8000 subjects, 130 questions). However, even after many experiments, it still takes 14 hrs for Rstan to run 1000 iterations on a subset of the data (1000 subjects, 80 questions). Then, it may take days to run the model on the whole dataset.

However, due to the time limit imposed by the computing cluster I run on, it is not feasible to let the model run for days. Hence, I wonder whether I can continue warmup/sampling after interruption. For example, is it possible to save the model as an R object after 1000 iterations and call the object later to continue warmup / sampling for another 1000 iterations?

Thank you very much!

1 Like
2

This can be done with CmdStan, CmdStanPy (beta on github) and PyStan (current develop on github) interfaces.

Basically you need to run warmup n-iters and and then extract stepsize, inverse metric and location of last warm-up.

I added a small example how to continue sampling with PyStan (so you would need to edit that example to handle only warmup).

Ps. There is currently a small bug with PyStan, so it needs at least 1 sample iter

3

Thank you very much! I will take a look at PyStan!

4

There was a small bug with stepsize and adapt_engaged needs to be False. These are now fixed in PyStan master branch.

I created a small example which shows how to use the inv_metric option. This same can be done with CmdStan (e.g. using CmdStanPy).

5

That was very helpful. I work in R and can’t find where to pass the inverse metric to the new object. Any help on that front is much-appreciated.

Here’s some example code that gets most of the way there.

library(stringr)
library(ggplot2)
library(rstan)

# Functions to get the last samples and adaptation information
extract_final_sample <- function(samples){
  draws <- extract(samples, permuted = FALSE, inc_warmup = TRUE)
  num_draws <- dim(draws)[1]

  list(as.list(draws[num_draws, , ]))
}

create_adaptation_list <- function(samples){
  adapt_string <- get_adaptation_info(samples)[[1]]
  sample_stepsize <- str_extract(adapt_string, '(?<=Step size = )[0-9\\.]+')
  sample_invmetric <- str_extract(adapt_string, '(?<=Diagonal elements of inverse mass matrix:\n# )([0-9\\.]+, )+[0-9\\.]+')

  list(adapt_engaged = FALSE,
       stepsize = as.numeric(sample_stepsize))

  #list(adapt_engaged = FALSE,
  #    stepsize = as.numeric(sample_stepsize),
  #    inv_metric = as.numeric(unlist(str_split(sample_invmetric, ', '))))
}

# Create data and simple model
num_obs <- 100
model_list <- list(N = num_obs,
                   x = rnorm(num_obs))

model_obj <- stan_model(model_code = '
                        data{
                          int<lower = 0> N;
                          vector[N] x;
                        }
                        parameters{
                          real mu;
                          real<lower = 0> sigma;
                        }
                        model{
                        x ~ normal(mu, sigma);
                        }
                        ')

# Draw initial sample with warmup samples
samples <- sampling(model_obj, model_list, iter = 250, warmup = 200, chains = 1)

# Try to draw samples starting from previous values
samples_cont <- sampling(model_obj, model_list, chains = 1, iter = 250, warmup = 0,
                         init = extract_final_sample(samples),
                         control = create_adaptation_list(samples))

# Plot the draws in sequence
mu_samples <- c(extract(samples, par = 'mu', permuted = FALSE, inc_warmup = TRUE),
             extract(samples_cont, par = 'mu', permuted = FALSE, inc_warmup = TRUE))

sample_data <- data.frame(mu = mu_samples,
                             id = ordered(rep(1:2, each = 250)),
                             num = 1:500)

ggplot(sample_data, aes(num, mu, color = id)) +
  geom_line() +
  theme_bw()
6

I think it is not yet implemented (I think @bgoodri could have some help with this feature).

Also RStan would need some way to combine draws. I used ArviZ on purpose. I think I could wrap StanFit4Model on PyStan to combine fits, but it was easier to work xarray Dataset / InferenceData.

7

there’s always CmdStan - you create a file in Rdump or JSON format which contains variable inv_metric and use argument metric_file=<filename>

8

Thank you. I’ll look into adapting it to CmdStan.

9

this is a common request - the question is do you want to continue warmup or do you want to continue sampling?

to continue sampling, specifying stepsize and metric_file with arguments adapt engaged=0 (in CmdStan, see page 64 of current CmdStan manual https://github.com/stan-dev/cmdstan/releases/download/v2.19.1/cmdstan-guide-2.19.1.pdf to figure out which argument comes in which order)

to continue warmup, specifying stepsize and metric will set the starting point for adaptation - see the section in the Stan manual on how warmup works -

1 Like
10

after reading the post, it looks like the cmdstanr backend for brms would be recommended. Then the inv_metric parameter can be extracted, which looks like it is the hard one to get your hands on if you want to avoid a burn in process. cmdstanr backend for brms is relatively new, but maybe there is some example code that someone has that performs the restart process from where a previous brms sampling exercise left off?

1 Like
11

@Cynon, with pure .R script you may get values of inv_matrix with cmdstanr functions, something like the following.
With given:

path is full path/to/your/csv files (stan output)
folder is “some_folder”
i is some chain number, i.e. 1 (if there are more chains you simply add a loop)

csv_stan <- list.files(path = path, pattern = "*.csv", full.names = T)
cmdstan_obj <- cmdstanr::read_cmdstan_csv(csv_stan)

data_metric <- list("inv_metric"=cmdstan_obj$inv_metric[[i]])
file_metric <- paste0(folder, "/metric_chain_", toString(i), ".json")
cmdstanr::write_stan_json(data = data_metric, file = file_metric)

Then you will have metric_chain_1.json file that can be used for new chain run.