profiling with pystan and speeding things up

[Ian Carroll]

The FAQ has some guidance for RStan users on profiling.

Is there a similar tool for PyStan users?

The Stan model I'd like to profile is attached, along with pickled sample of the data and a python script to compile and run.

Along with advice on how to profile it myself, expert opinions on improving efficiency are very welcome!

Here's the model in case the attachment fails:

data {

    int N; // number of samples

    int K; // number of intercepts

    int L; // number of categoricals predictors

    int C[L]; // position of last coefficient in each category

    int M; // number of non-zero values in sparse design matrix for intercepts

    vector[M] c_w;   // design matrix in compressed row storage where c_w[i]

    int c_v[M];      // is a value from column c_v[i], row j values begin

    int c_u[N + 1];  // at position c_u[j] in c_w, and c_u[N + 1] = M.

    int P; // number of covariates

    matrix[N, P] x; // covariates

    int R; // number of non-zero rows in sample matrix augmented with totals

    int W; // number of non-zero values in sparse matrix A

    //vector[W] A_w;  // A in compressed row storage where every A_w[i]

    int A_v[W];      // is one, it's from column A_v[i], row j values

    int A_u[R + 1];  // begin at position A_u[j] in A_w, and A_u[R + 1] = W

    int Q; // number of non-zero values in sparse sample matrix

    vector[Q] y;   // response value, from sparse matrix,

    int y_i[Q];    // such that y[k] reports have more than

    vector[Q] y_j; // y_j[k] animals from market:county y_i[k]

}

parameters {

    vector[K] alpha;

    vector[P] beta;

    real<lower=0> sigma; // SD for market-level random effect

}

transformed parameters {

    vector[N] log_a;

    log_a <- csr_matrix_times_vector(N, K, c_w, c_v, c_u, alpha) + x * beta;

}

model {

    // Priors

    alpha[C[1]] ~ normal(0, 10);

    alpha[(C[2] + 1):C[3]]  ~ normal(0, 10);

    beta ~ normal(0, 10);

    sigma ~ cauchy(0, 2.5);

    // Hierarchies

    alpha[(C[1] + 1):C[2]] ~ normal(0, sigma);

    // Likelihood

    {

vector[N] a;

vector[R] b;

int j;

a <- exp(log_a);

b <- rep_vector(0, R);

j <- 1;

// implement csr_matrix_times_vector for binary matrix

        // doesn't do any better or worse than the function, though

for (i in 1:R) {

   while (j < A_u[i + 1]) {

b[i] <- b[i] +  a[A_v[j]];

j <- j + 1;

   }

}

increment_log_prob(dot_product(log(b[y_i] + y_j), y));

    }

}

[Bob Carpenter]

The actual sampling for Stan programs doesn't vary by
platform (other than that the platforms other than CmdStan
save the sample in memory) --- they all use the same C++
sampler on the back end.

There's basically zero cost for a function call (relative
to gradient calculations, which is where almost all the work
in Stan goes).

Depending on how much data you have, you might want to
use the non-centered parameterization:

alpha_raw ~ normal(0, 1);

and then rescale to get alpha.

You can also make log_a a local variable in the model if you don't
need to save it.

You can probably just define b with the appropriately indexed
sum(), which is more efficient than repeated addition (fewer nodes
in the expression graph, which is what the derivatives run on).

- Bob

> --
> You received this message because you are subscribed to the Google Groups "Stan users mailing list" group.
> To unsubscribe from this group and stop receiving emails from it, send an email to stan-users+...@googlegroups.com.
> To post to this group, send email to stan-...@googlegroups.com.
> For more options, visit https://groups.google.com/d/optout.
> <example.py><model.stan><data.pickle>

[Andrew Gelman]

Is this really true? I thought that with autodiff, a gradient calculation might just cost something like 3x the cost of a function call.

Or is your point that Stan calculates gradients much more often than it does function calls?

I guess I’m just confused (as usual).

A

[Michael Betancourt]

When we compute the gradients we get the function evaluation for free,
and given that we’re always computing gradients we almost always have
the function value around.

[Bob Carpenter]

And no, it's not 3x the cost of a function call. All those
people at the DoE talk who told you otherwise had never
profiled anything.

And it depends on the function. See, for example, this neat arXiv paper,
which has lots of graphs (though you won't like the tick
marks, the grey backgrounds, the legends not labeling the lines,
etc. etc. etc.):

Bob Carpenter, Matthew D. Hoffman, Marcus Brubaker, Daniel Lee, Peter Li, Michael Betancourt.
2015.
The Stan Math Library: Reverse-Mode Automatic Differentiation in C++.
http://arxiv.org/abs/1509.07164

You can see (p. 68) that sum is a factor of 16 and product a factor of 4, but
they're worse for smaller expressions. Our vectorized sum is a lot faster.

- Bob

P.S. Here are those plots, all relative to Stan, including the double value
that just evaluates the function itself.

[Andrew Gelman]

Wow—this paper is great. But, yes, I think the graphs could be done much better. It’s too hard to try to decode the colors and go back and forth between the labels.

> --
> You received this message because you are subscribed to the Google Groups "Stan users mailing list" group.
> To unsubscribe from this group and stop receiving emails from it, send an email to stan-users+...@googlegroups.com.
> To post to this group, send email to stan-...@googlegroups.com.
> For more options, visit https://groups.google.com/d/optout.

> <product_rel_eval.pdf><sum_rel_eval.pdf>

[Ian Carroll]

You guys are funny. I like the enthusiasm.

I guess I'll add profiling to the PyStan feature request then, as I'm not so keen to set up shop in R for that one task.

Bob - Thanks for the suggestions, I think the implementation below is about twice as fast, but I'm not sure I caught your drift with non-centering. Would you mind checking at the lines marked with ** below?

New Stan model definition (new data also attached):

data {

    int N;  // number of samples

    int K;  // number of intercepts

    int L;  // number of categorical predictors

    int C[L];  // position of last coefficient in each category

    int M;  // number of non-zero values in sparse design matrix for intercepts

    vector[M] c_w;   // design matrix in compressed row storage where c_w[i]

    int c_v[M];      // is a value from column c_v[i], row j values begin

    int c_u[N + 1];  // at position c_u[j] in c_w, and c_u[N + 1] = M.

    int P;  // number of covariates

    matrix[N, P] x;  // covariates

    int R_1;

    int R_2;

    int b_1[R_1];

    int b_1_idx[R_1];

    int b_2[R_2 + 1];

    int b_2_idx[R_2];

    int Q;  // number of non-zero values in sparse sample matrix

    vector[Q] y;    // response value, from sparse matrix,

    int y_i[Q];     // such that y[k] reports have more than

    vector[Q] y_j;  // y_j[k] animals from market:county y_i[k]

}

parameters {

    vector[C[1]] alpha_1;

    vector[C[2] - C[1]] alpha_2_raw;

    vector[C[3] - C[2]] alpha_3;

    vector[P] beta;

    real<lower=0> sigma;  // SD for market-level random effect

}

model {

    // Prior

    append_row(alpha_1, alpha_3) ~ normal(0, 100);

    alpha_2_raw ~ normal(0, 1); // **

    sigma ~ cauchy(0, 2.5);

    beta ~ normal(0, 100);

    // Likelihood

    {

vector[K] alpha;

vector[N] a;

vector[R_1 + R_2] b;

alpha <- append_row(append_row(alpha_1, alpha_2_raw * sigma), alpha_3);     // **

a <- csr_matrix_times_vector(N, K, c_w, c_v, c_u, alpha) + x * beta;

a <- exp(a);

b[b_1_idx] <- a[b_1];

for (i in 1:R_2) {

   b[b_2_idx[i]] <- sum(a[b_2[i]:(b_2[i+1] - 1)]);

}

increment_log_prob(dot_product(log(b[y_i] + y_j), y));

    }

}

[Bob Carpenter]

Thanks for asking the question with pointers to what
I should look at.

Yes, that's the right way to do non-centering for alpha_2.

You could do the same for alpha_1 and beta (multiply a unit normal
by 100). Generally, we recommend tighter priors. So unless you
think 200 or -200 is a resaonable value for elements of alpha_1
or beta, that prior's too fat.

- Bob

> <data.pickle>