Defining a Pyro model

Hello, I read over the Pyro documentation and I am thinking perhaps the way I defined my Bayesian model is incorrect? Could you take a look at the summerized version of my code below and see if they are ok? Thank you very much for your help.

# convert `myFrequentistModel` into a pyro module.
module.to_pyro_module_(myFrequentistModel)

# convert the uppermost (the 24th) layer of `myFrequentistModel` into a
# Bayesian layer.
for m in myFrequentistModel.layer[23].modules():
     for name, value in list(m.named_parameters(recurse=False)):
         if name != "_dummy_param":
            setattr(m, name, module.PyroSample(prior=dist.Normal(0, 1)
                                              .expand(value.shape)
                                              .to_event(value.dim())))

# define likelihood function for our Bayesian layer.
class MyModel(PyroModule):
    
    def __init__(self,  myFrequentistModel, name=""):
        self._pyro_name = name
        self._pyro_context = pyro.nn.module._Context()
        self._pyro_params = myFrequentistModel.parameters()
        self._modules = myFrequentistModel.modules()
        super(MyModel, self).__init__()

    def forward(self, myFrequentistModel, input, yLabel):

        softmax_tensor = myFrequentistModel(input)
  
        pyro.sample("y",
                    dist.Multinomial(1, probs = softmax_tensor),
                    obs = yLabel)

        return softmax_tensor

myPyroModel = MyModel(myFrequentistModel)

# train myPyroModel...
trainedPyroModel = train(myPyroModel)

# turn on the evaluation mode.
trainedPyroModel.eval()
pred_obj = Predictive(trainedPyroModel, guide=guide,
                                  num_samples = 100)
                                  
# make predictions on the test set
predicted_softmax_tensor = pred_obj.call(myFrequentistModel, test_input, yLabel=None)

Hello,
Would it make more sense to change the forward function of MyModel class like the following?:

# define likelihood function for our Bayesian layer.
class MyModel(PyroModule):
    
    def __init__(self,  myFrequentistModel, name=""):
        self._pyro_name = name
        self._pyro_context = pyro.nn.module._Context()
        self._pyro_params = myFrequentistModel.parameters()
        self._modules = myFrequentistModel.modules()
        super(MyModel, self).__init__()

    def forward(self, myFrequentistModel, input, yLabel, y = None):

        softmax_tensor = myFrequentistModel(input)
  
        return pyro.sample("y",
                    dist.Multinomial(1, probs = softmax_tensor),
                    obs = yLabel)

Thank you,

Hi @h56cho, you shouldn’t need to set private attributes ._* directly. How about this simplification:

class MyModel(PyroModule):
    def __init__(self,  myFrequentistModel):
        super().__init__()
        self.model = myFrequentistModel  # everything should happen automatically here

    def forward(self, input, yLabel):
        softmax_tensor = self.model(input)
        return pyro.sample("y",
                    dist.Multinomial(1, probs = softmax_tensor),
                    obs = yLabel)