Attribute Error when setting number of chains greater than 1 in MCMC

Misc.
1

I have defined a very simple linear regression model where I am using HMC to sample the coefficients. However, when setting num_chains > 1 in the MCMC class, I get an AttributeError.

Software and Hardware:
python = 3.10.6
pyro = 1.8.2+2e3bd02
Laptop = Apple M1 Pro running macOS Monterey

First, the data used for the model:

try: from palmerpenguins import load_penguins
except ImportError:
    from pip._internal import main as pip
    pip(['install', '--user', 'palmerpenguins'])
    from palmerpenguins import load_penguins

penguins = load_penguins()
penguins.dropna(how='any', axis=0, inplace=True)
adelie_mask = (penguins['species'] == 'Adelie')
adelie_flipper_length = torch.from_numpy(penguins.loc[adelie_mask, 'flipper_length_mm'].values)
adelie_mass = torch.from_numpy(penguins.loc[adelie_mask, 'body_mass_g'].values)

Where the model and MCMC are:

def linear_model(flipper_length, mass=None):

    sigma = pyro.sample('sigma', dist.HalfNormal(2000.))
    beta_0 = pyro.sample('beta_0', dist.Normal(0., 4000.))
    beta_1 = pyro.sample('beta_1', dist.Normal(0., 4000.))
    mu = pyro.deterministic('mu', beta_0 + beta_1 * flipper_length)

    with pyro.plate('plate'):   
        preds = pyro.sample('mass', dist.Normal(mu, sigma), obs=mass)  

kernel = NUTS(linear_model, adapt_step_size=True)
mcmc_simple = MCMC(kernel, num_samples=500, warmup_steps=300, num_chains=4)
mcmc_simple.run(flipper_length=adelie_flipper_length, mass=adelie_mass)

When setting num_chains > 1 in MCMC(), I get an Attribute Error. Interestingly, the number of Tracebacks or AttributeErrors is equal to the number of chains specified. In the example above, num_chains=4, which amounts to four AttributeErrors:

Traceback (most recent call last):
Traceback (most recent call last):
Traceback (most recent call last):
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "<string>", line 1, in <module>
  File "<string>", line 1, in <module>
  File "<string>", line 1, in <module>
  File "/Users/gabestechschulte/miniforge3/envs/probs/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main
  File "/Users/gabestechschulte/miniforge3/envs/probs/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main
  File "/Users/gabestechschulte/miniforge3/envs/probs/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main
  File "/Users/gabestechschulte/miniforge3/envs/probs/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main
        exitcode = _main(fd, parent_sentinel)exitcode = _main(fd, parent_sentinel)

  File "/Users/gabestechschulte/miniforge3/envs/probs/lib/python3.10/multiprocessing/spawn.py", line 126, in _main
  File "/Users/gabestechschulte/miniforge3/envs/probs/lib/python3.10/multiprocessing/spawn.py", line 126, in _main
    exitcode = _main(fd, parent_sentinel)    
exitcode = _main(fd, parent_sentinel)
  File "/Users/gabestechschulte/miniforge3/envs/probs/lib/python3.10/multiprocessing/spawn.py", line 126, in _main
  File "/Users/gabestechschulte/miniforge3/envs/probs/lib/python3.10/multiprocessing/spawn.py", line 126, in _main
    self = reduction.pickle.load(from_parent)
AttributeError:     Can't get attribute 'linear_model' on <module '__main__' (built-in)>self = reduction.pickle.load(from_parent)

AttributeError: Can't get attribute 'linear_model' on <module '__main__' (built-in)>
    self = reduction.pickle.load(from_parent)
AttributeError: Can't get attribute 'linear_model' on <module '__main__' (built-in)>
    self = reduction.pickle.load(from_parent)
AttributeError: Can't get attribute 'linear_model' on <module '__main__' (built-in)>

This error does not seem to be similar to this issue involving Python 3.10 and PyTorch since my issue seems to stemming from the multiprocessing/spawn.py file. Has anyone also encountered this issue? Any help would be greatly appreciated.

Thank you for your time :smile:

2

Hey gang, just a quick update on this before I dive deeper into the bug. It seems that when defining the model as a function and running the MCMC sampler in a Juptyter Notebook, the spawn.py file is never assigned the attribute ‘linear_model’.

However, if you define and run the model in the main program (source file), i.e., if __name__ == "__main__", then num_chains > 1 works as normal.

Therefore, this problem seems to be with Jupyter Notebook. Current solution is to develop Pyro models in a main module.

Code:

import torch
import pyro.distributions as dist
import pyro
from pyro.infer import Predictive, NUTS, MCMC
from palmerpenguins import load_penguins

def linear_model(flipper_length, mass=None):

    sigma = pyro.sample('sigma', dist.HalfNormal(2000.))
    beta_0 = pyro.sample('beta_0', dist.Normal(0., 4000.))
    beta_1 = pyro.sample('beta_1', dist.Normal(0., 4000.))
    mu = pyro.deterministic('mu', beta_0 + beta_1 * flipper_length)

    with pyro.plate('plate'):   
        preds = pyro.sample('mass', dist.Normal(mu, sigma), obs=mass)  

def main():
    
    penguins = load_penguins()
    penguins.dropna(how='any', axis=0, inplace=True)
    adelie_mask = (penguins['species'] == 'Adelie')
    adelie_flipper_length = torch.from_numpy(penguins.loc[adelie_mask, 'flipper_length_mm'].values)
    adelie_mass = torch.from_numpy(penguins.loc[adelie_mask, 'body_mass_g'].values)

    kernel = NUTS(linear_model, adapt_step_size=True)
    mcmc_simple = MCMC(kernel, num_samples=500, warmup_steps=300, num_chains=4)
    mcmc_simple.run(flipper_length=adelie_flipper_length, mass=adelie_mass)

if __name__ == '__main__':
    main()