User-defined distribution for SVI log_prob method

I’m trying to define my own distribution with its own sample and log_prob methods. The distribution is supposed to model the noise in an observed k-mer (string of nucleotides A,C,G,T). After importing the class, I can use pyro.sample(“obs”, NoiseModel(x))) to generate a noisy version of k-mer x (e.g., if x=ACGT then I might sample ACGA). However, when I define my pyro model and try to do inference with SVI, I get the error AttributeError: ‘function’ object has no attribute ‘log_prob’. I think it might have to do with the fact that in the second case I am passing the optional argument data. However, I am not sure how this argument is messing with my log_prob method. Thanks in advance!

###############################################################
# Class definition (mymodule.py)
###############################################################
import numpy as np
from pyro.distributions import Distribution

# Constants
NUCSET = set('ACGT')

# Noise model for k-mers
class NoiseModel(Distribution):

    def __init__(self,bin_mean_kmer):
        self.bin_mean_kmer = bin_mean_kmer

    def sample(self,sample_shape=torch.Size()):
        noisy_kmer = ''
        kmer = dec_seq(self.bin_mean_kmer)
        for nuc in kmer:
            if np.random.sample()<=0.8:
                noisy_kmer += nuc
            else:
                noisy_kmer += np.random.choice(tuple(NUCSET.difference(set(nuc))))

        return noisy_kmer

    def log_prob(self,value):

        log_p = 0.0
        kmer = dec_seq(self.bin_mean_kmer)
        for i in range(len(kmer)):
            log_p += np.log(0.8) if kmer[i]==value[i] else np.log(0.2)

        return log_p

# Decode k-mer
def dec_seq(x):

    y = ''
    for i in range(0,len(x),2):
        if x[i]==0 and x[i+1]==0:
            y = y + 'A'
        elif x[i]==0 and x[i+1]==1:
            y = y + 'C'
        elif x[i]==1 and x[i+1]==0:
            y = y + 'G'
        else:
            y = y + 'T'

    return y

# Encode k-mer
def enc_seq(y):

    x = [0] * (2*len(y))
    for i in range(len(y)):
        j = 2 * i
        if y[i]=='A':
            x[j]=0
            x[j+1]=0
        elif y[i]=='C':
            x[j]=0
            x[j+1]=1
        elif y[i]=='G':
            x[j]=1
            x[j+1]=0
        else:
            x[j]=1
            x[j+1]=1

    return x

###############################################################
# Pyro model
###############################################################

from mymodule import NoiseModel,enc_seq,dec_seq

# Sample binary version of k-mer
def sample_bin_rep():

    x = [0] * (2*K)
    for i in range(2*K):
        x[i] = 1 if np.random.uniform()>0.5 else 0

    return torch.tensor(x)

# Stick-breaking function
def mix_weights(beta):
    beta1m_cumprod = (1 - beta).cumprod(-1)
    return F.pad(beta, (0, 1), value=1) * F.pad(beta1m_cumprod, (1, 0), value=1)

# Define pyro model
def model(data):
    with pyro.plate("beta_plate", T-1):
        beta = pyro.sample("beta", Beta(1, alpha))

    with pyro.plate("true_flank_plate", T):
        true_flank = pyro.sample("true_flank", sample_bin_rep)

    with pyro.plate("data", N):
        z = pyro.sample("z", Categorical(mix_weights(beta)))
        pyro.sample("obs", NoiseModel(true_flank[z]), obs=data)

Hi @jordi, it’s tough to answer without seeing your error message and what that function object is, but here are some tips on creating custom distributions:

  1. Use pytorch rather than numpy everywhere.
  2. Inherit from TorchDistribution rather than Distribution.
  3. Remember to call super().__init__()
  4. Add metadata like support, arg_constraints, batch_shape, and event_shape.

I’d recommend taking a look at Eli Weinstein’s sequence distributions or other distribution implementations for example code. We plan to make a how-to-create-a-distribution tutorial, but things are busy; contributions welcome :slightly_smiling_face:

Hi @fritzo, thank you so much for your quick response. I included points 1 and 2, but I am not sure how to deal with points 3 and 4 since I am relatively new, but I will look into it. I’ve been checking implementation examples as you suggest but like I said, this is all a bit foreign for me, so I think a how-to-create-a-distribution tutorial would make it much more accessible. In fact, I found all the tutorials in the website extremely helpful. I’m actually basing this code off of the DPMM examples you all have. In any case, here is the output I get after incorporating suggestions 1 and 2:

>>> K = 10          # K-mer size
>>> T = 10          # Truncation to T components
>>> N = 20          # Number of observations
>>> z,data = sim_model()
>>> alpha = 1.1
>>> n_iter = 1500
>>> optim = Adam({"lr": 0.05})
>>> svi = SVI(model, guide, optim, loss=Trace_ELBO())
>>> losses = []
>>> train(n_iter,data)
  0%|                                                                                                                                                                                                                                  | 0/1500 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 4, in train
  File "/home/users/jabante/python-env/metagenomics_env/lib/python3.6/site-packages/pyro/infer/svi.py", line 145, in step
    loss = self.loss_and_grads(self.model, self.guide, *args, **kwargs)
  File "/home/users/jabante/python-env/metagenomics_env/lib/python3.6/site-packages/pyro/infer/trace_elbo.py", line 140, in loss_and_grads
    for model_trace, guide_trace in self._get_traces(model, guide, args, kwargs):
  File "/home/users/jabante/python-env/metagenomics_env/lib/python3.6/site-packages/pyro/infer/elbo.py", line 186, in _get_traces
    yield self._get_trace(model, guide, args, kwargs)
  File "/home/users/jabante/python-env/metagenomics_env/lib/python3.6/site-packages/pyro/infer/trace_elbo.py", line 58, in _get_trace
    "flat", self.max_plate_nesting, model, guide, args, kwargs
  File "/home/users/jabante/python-env/metagenomics_env/lib/python3.6/site-packages/pyro/infer/enum.py", line 64, in get_importance_trace
    model_trace.compute_log_prob()
  File "/home/users/jabante/python-env/metagenomics_env/lib/python3.6/site-packages/pyro/poutine/trace_struct.py", line 230, in compute_log_prob
    log_p = site["fn"].log_prob(
AttributeError: 'function' object has no attribute 'log_prob'

Hi @jordi it looks like you’re error is with the line

true_flank = pyro.sample("true_flank", sample_bin_rep)

which you can verify by printing site["name"] or site["fn"] in a debugger. You’ll need to replace that function with an actual distribution object, e.g.

true_flank = pyro.sample("true_flank", dist.Categorical(logits=torch.zeros(2 * K, 2)))

Actually I don’t understand why that’s in the model. It looks like it should be in a data generation step outside of the model, and you’d pass in true_flank as an argument to the model.

Thanks for checking the code, @fritzo. I see, I was too focused on the other plate that I didn’t even realize the issue was with true_flank_plate. I need true_flank_plate though because that’s the unobserved layer I’m after. I just get to observe noisy versions of true flank (true k-mer) and my goal is to try to infer from the data which true k-mer generated the observed k-mers, and I’m modeling the former as a Dirichlet Process with a measure defined over the set of all 4^k k-mers since when k is large the state space is practically infinite in size. For efficient sampling though, I am not using a Categorical with 4^k categories. Instead, I assume independence of each nucleotide, represent each one with two binary numbers, and sample from a Bernoulli a vector of size 2k.

After struggling for a while with the user defined distribution NoiseModel I got it to work (I think). I would be happy to contribute by sharing this if you all want to use it as an example of a user defined distribution. I’m sure it can be improved substantially though, so I would very much welcome your edits. How should I proceed to share it with you all?

Hi @jordi, we’d love help making a distribution tutorial, and your distribution looks like a good starting point. If you’d like to write the first version of a tutorial yourself, I’d recommend forking an existing tutorial such as the tensor shapes tutorial (rendered html | source notebook). That may involve a lot of boilerplate; if you want to collaborate on a PR I guess you could put up a gist with your distribution, then I could create a notebook draft with your gist, and you and I could add text to the notebook. Either way thanks for contributing! :grinning_face_with_smiling_eyes: