How can I speed up the execution of my model?

Questions
1

Hi,
I am trying to use PyMC3 to do a Bayesian fitting of diffusion MRI signal. For those who are not “fluent” in MRI, the principle is to probe the diffusion of water in an organ of interest. To do so, a magnetic field gradient is applied to the organ of interest defore the signal is measured. A simple mono exponential decay model links the signal measured in each voxel (a 3D pixel) to the strength of the applied gradient following a simple mono-exponential decay:

S(b)=S_0\exp\{-b\times ADC\}

Where b is the magnetic field gradient strength, called b-value, ADC is the apparent diffusion coefficient, and S_0 is the signal with no gradient applied.
The b-values are set by the experimental design and are the same for every voxels.
The aim is to estimate ADC and to a lesser extent S_0.
We assume a Gaussian noise with standard deviation \sigma.
I wanted to use PyMC3 to implement a very simple model, that is defined, for every voxels i and b-value j, as:

y_{ij}=S_{0_i}\exp\{-b_j\times ADC_i\} + \epsilon_{ij}
\epsilon_{ij} \sim \mathcal{N}(0,\sigma_i)
\sigma_i\sim \mathrm{HalfCauchy}(0.5)
S_{0_i}\sim \mathcal{N}(0,1e5)
ADCi_\sim \mathrm{Log-}\mathcal{N}(4e-3,1)

I know the model is quite naive and simple: every voxel is independent from the other, the prior are fixed and non informative, but I chose it as a training in PyMC3, before moving to more complex modeling.
Here is the code I wrote:

Imports

%matplotlib inline
import numpy as np
import pymc3 as pm
import matplotlib.pyplot as plt
import seaborn as sns
plt.style.use('seaborn-darkgrid')

Data generation

M=5 # number of voxel in my experiment
bvalues=np.array([0,10,20,50,100,200,500,1000],dtype=np.float) # vector of b-values 
N=bvalues.size # number of bvalue

ADC=2e-3 # Same ADC and S0 for every voxels
S0=1
Sb=S0*np.exp(-bvalues*ADC) # noise free signal
sigma=0.1
data=np.random.normal(loc=Sb,scale=sigma, size=(M,N)) # Add noise and create the signal for M voxels

Here is an example of the MRI signal, as a function of the b-value

image
image.png1505×780 105 KB

PyMC3 model

with  pm.Model() as unpooled_model:
    adc=pm.Lognormal('adc',np.log(4e-3),sigma=1, shape=(M,1)) #setting shape=M does not allow the broadcasting to work properly
    s0=pm.Normal('s0',0,1e3, shape=(M,1)) 
    sigma = pm.HalfCauchy('sigma', .5,shape=(M,1))

    sb_est=s0*np.exp(-bvalues*adc)
    
    sb_like=pm.Normal('sb_like', mu=sb_est, sigma=sigma, observed=data )

The graph of the model is then:

image
image.png1632×592 58.2 KB

Inference

with unpooled_model:
    unpooled_trace = pm.sample(2000, tune=1000)

pm.traceplot(unpooled_trace, 
             var_names=['adc','s0','sigma'], 
             coords={'adc_dim_0': 0,'s0_dim_0': 0,'sigma_dim_0': 0});

image
image.jpg2175×1087 308 KB

At this point I am quite satisfied with the result. However, it took almost 30s for the algorithm to converge, which disappointed me of course.

So here is my simple question:

Do you have an idea to improve the computation time? I don’t know how to profile this kind of code, so I am quite helpless.
Do you think that implementing the monoexpontial model in theano could speed it up? Then how should it be done?

Thanks for the help!

2

Looking at the model this is already pretty well implemented, what you can consider doing is to change the LogNormal and the HalfCauchy prior to HalfNormal, which is usually a bit easier to sample from.

3

Ok.
Do you think that implementing the model with theano could speed up the code?
I don’t know exactly when theano is used in the background or when it is required to use it in my own code. I don’t know if using the theano exp function instead of the numpy function would improve the computation time, or would change anything.
I understood that np.exp() will use theano.tensor.exp() in the background, is that correct? In that case I guess that there is no benefit to use theano explicitely.

4

I think so too, but it is pretty straightforward to switch to theano version so I suggest you to go for it.

5

Disclaimer: I am not an expert here, but my two cents. Why is 30 s disappointing? Are you trying to scale to larger dataset? I read in the pymc docs that often pymc is equal to or slower than other packages bayesian packages when the data is small, as there is a lot of overhead compiling. But as you fit larger data the compiling overhead becomes a smaller percentage of compute time and thus pymc becomes faster than some of the other bayesian software packages. So if you need to scale the model it might not perform as badly as you think.

1 Like