Ch0ronomato
commented
2 days ago The script I used for reference
import arviz as az
import numpy as np
import multiprocessing
import pandas as pd
import pymc as pm
import pytensor as pt
import pytensor.tensor.random as ptr
import time
def main():
# Load the radon dataset
data = pd.read_csv(pm.get_data("radon.csv"))
data["log_radon"] = data["log_radon"].astype(np.float64)
county_idx, counties = pd.factorize(data.county)
coords = {"county": counties, "obs_id": np.arange(len(county_idx))}
# Create a simple hierarchical model for the radon dataset
with pm.Model(coords=coords, check_bounds=False) as model:
intercept = pm.Normal("intercept", sigma=10)
# County effects
raw = pm.ZeroSumNormal("county_raw", dims="county")
sd = pm.HalfNormal("county_sd")
county_effect = pm.Deterministic("county_effect", raw * sd, dims="county")
# Global floor effect
floor_effect = pm.Normal("floor_effect", sigma=2)
# County:floor interaction
raw = pm.ZeroSumNormal("county_floor_raw", dims="county")
sd = pm.HalfNormal("county_floor_sd")
county_floor_effect = pm.Deterministic(
"county_floor_effect", raw * sd, dims="county"
)
mu = (
intercept
+ county_effect[county_idx]
+ floor_effect * data.floor.values
+ county_floor_effect[county_idx] * data.floor.values
)
sigma = pm.HalfNormal("sigma", sigma=1.5)
pm.Normal(
"log_radon", mu=mu, sigma=sigma, observed=data.log_radon.values, dims="obs_id"
)
from pymc.model.transform.optimization import freeze_dims_and_data
model = freeze_dims_and_data(model)
for mode in ("NUMBA", "PYTORCH"):
start = time.perf_counter()
trace = pm.sample(
model=model,
cores=1,
chains=1,
tune=500,
draws=500,
progressbar=False,
compute_convergence_checks=False,
return_inferencedata=False,
compile_kwargs=dict(mode=mode)
)
end = time.perf_counter()
idata = pm.to_inference_data(trace, model=model)
print(az.summary(idata, kind="diagnostics"))
print(mode, trace._report.t_sampling, end - start)
if __name__ == "__main__":
multiprocessing.freeze_support()
main()
ricardoV94
Description
@ricardoV94 did a nice perf improvement in https://github.com/pymc-devs/pymc/pull/7578 to try to speedup jitted backends. I tried out torch as well. The model performed quite slow.
We need to investigate why
When doing perf evaluations, keep in mind that torch does a lot of caching. If you want a truly cache-less eval, you can either add
torch.compiler.reset()or set the env variable to disable the dynamo cache (google it).