PyroModule, PyroSample, PyroParam and GRU

Hi!

I am writing again about the same issue posted here: PyroModule for LSTM.
I have been following this example : http://docs.pyro.ai/en/stable/nn.html#pyro.nn.module.to_pyro_module_

But either there is a bug or I am not doing something in the right order, hehe :smiley:

Pseudo code of my model:

def Convert_to_PyroSample(model):
    for m in model.modules():
        for name, value in list(m.named_parameters(recurse=False)):
            setattr(m, name, PyroSample(prior=dist.Normal(0, 1).expand(value.shape).to_event(value.dim())))

class GRU(nn.Module): # I am not sure whether this should already be PyroModule (with either option the error persists)
        def __init__(self):
            super(...)
            self.GRU = nn.GRU(flags...)
            self.h_0   = nn.Parameter(torch.rand(GRU_hidden_size), requires_grad = True) #Shuld it be converted to PyroParam here?
       def forward(self,input):
             to_pyro_module(self.GRU)
             Convert_to_PyroSample(self.GRU)
             h_0_contig = PyroParam(self.h_0.repeat(...).contiguous())   # Not convinced about this
             output,_ = self.GRU(input, h_0_contig)
             return output

a) If I try to convert h_0_contig to PyroParam I get this error:

File “/home/…/anaconda3/lib/python3.7/site-packages/torch/nn/modules/rnn.py”, line 175, in check_hidden_size
if hx.size() != expected_hidden_size:
AttributeError: ‘PyroParam’ object has no attribute ‘size’

b) If I don’t bother with PyroParam I get the same error as in https://forum.pyro.ai/t/pyromodule-for-lstm/1596.:

File “/home/…/anaconda3/lib/python3.7/site-packages/torch/nn/modules/rnn.py”, line 716, in forward
self.dropout, self.training, self.bidirectional, self.batch_first)
TypeError: expected Tensor as element 0 in argument 2, but got PyroSample

I am just confused on the semantics, thanks for your help in advance! I am using pyro 1.3.0

Best wishes

Hi @artistworking,
I think there’s just some confusion between our old pyro.param usage and the new PyroParam usage. Whereas pyro.param belongs in the .forward() method, PyroParam now belongs in the .__init__() method, and you access those params in the .forward() method. Here’s an attempt:

def Convert_to_PyroSample(model):
    for m in model.modules():
        for name, value in list(m.named_parameters(recurse=False)):
            setattr(m, name, PyroSample(prior=dist.Normal(0, 1).expand(value.shape).to_event(value.dim())))

class GRU(PyroModule):
    def __init__(self):
        super(...)
        self.GRU = nn.GRU(flags...)
        self.h_0 = PyroParam(torch.rand(GRU_hidden_size))
        Convert_to_PyroSample(self.GRU)
   def forward(self, input):
        # In the following line, the self.h_0 lookup now
        # triggers an internal pyro.param() call:
        h_0_contig = self.h_0.repeat(...).contiguous()
        output,_ = self.GRU(input, h_0_contig)
        return output

Let me know if that still doesn’t work. Also you might try grepping around the Pyro codebase for other internal uses of PyroModule, e.g. in https://github.com/pyro-ppl/pyro/blob/master/tests/nn/test_module.py

Thanks so much for your reply. I have switched the code as suggested and I also had a look at the test_module.py :slight_smile: . I can see I mixed some stuff up :upside_down_face:

I have made more changes actually, because otherwise I could not convert the GRU parameters to PyroSample. Convert_to_PyroSample was not working, see error below:

File “/home/…/Example.py”, line 61, in Convert_to_PyroSample
setattr(m, name, PyroSample(prior=dist.Normal(0, 1).expand(value.shape).to_event(value.dim())))
File “/home/…/anaconda3/lib/python3.7/site-packages/torch/nn/modules/rnn.py”, line 97, in setattr
super(RNNBase, self).setattr(attr, value)
File “/home/…/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 595, in setattr
.format(torch.typename(value), name))
TypeError: cannot assign ‘pyro.nn.module.PyroSample’ as parameter ‘weight_ih_l0’ (torch.nn.Parameter or None expected)

def Convert_to_PyroSample(model):
    for m in model.modules():
        for name, value in list(m.named_parameters(recurse=False)):
            setattr(m, name, PyroSample(prior=dist.Normal(0, 1).expand(value.shape).to_event(value.dim())))

class GRU(PyroModule):
    def __init__(self):
        super(...)
        self.GRU = PyroModule[nn.GRU](flags...) <--- Added that
        self.h_0 = PyroParam(torch.rand(GRU_hidden_size))
        Convert_to_PyroSample(self.GRU)
   def forward(self, input):
        # In the following line, the self.h_0 lookup now
        # triggers an internal pyro.param() call:
        h_0_contig = self.h_0.repeat(...).contiguous()
        assert isinstance(self.GRU, PyroModule) --> It is a PyroModule
        output,_ = self.GRU(input, h_0_contig)
        return output

However, the GRU does not seem to have been affected by the PyroModule. I get the same error if I manually assign each of the weights and bias of the GRU to PyroSample:

File “/home/…/Example.py”, line 85, in forward
output, _ = self.GRU(input,h_0_contig)
File “/home/…/anaconda3/lib/python3.7/site-packages/pyro/nn/module.py”, line 288, in call
return super().call(*args, **kwargs)
File “/home/…/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 532, in call
result = self.forward(*input, **kwargs)
File “/home/…/anaconda3/lib/python3.7/site-packages/torch/nn/modules/rnn.py”, line 716, in forward
self.dropout, self.training, self.bidirectional, self.batch_first)
TypeError: expected Tensor as element 0 in argument 2, but got PyroSample

Thanks again!

Hi @artistworking, I think this may be a low-level incompatibility between PyroModule and nn.RNN. Can you see if the following workaround helps?

class GRU(PyroModule):
    ...
    def forward(self, input):
        h_0_contig = self.h_0.repeat(...).contiguous()
        self.GRU._apply(lambda t: t)  # <--- recomputes GRU._flat_weights
        output,_ = self.GRU(input, h_0_contig)
        return output

For context, it appears that nn.RNN caches flat views of its parameters, and that the cache becomes invalid at each Pyro sample call, i.e. when new parameters are sampled in GRU.forward. I’m not sure how to fix, maybe we can provide a pyro.nn.PyroRNN wrapper or something. Feel free to file a bug at https://github.com/pyro-ppl/pyro/issues/new or I can do so if you prefer. EDIT I’ve filed an issue here: https://github.com/pyro-ppl/pyro/issues/2390