Bayesian Regression

class BayesianRegression(PyroModule):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = PyroModule[nn.Linear](in_features, out_features)
        self.linear.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2))
        self.linear.bias = PyroSample(dist.Normal(0., 1.).expand([out_features]).to_event(1))

    def forward(self, x, y=None):
        sigma = pyro.sample("sigma", dist.Uniform(0., 10.))
        mean = self.linear(x).squeeze(-1)
        with pyro.plate("data", x.shape[0]):
            obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
        return mean

Hi,
How do I modify this class so that it returns two outputs for every input?
What I mean is, I want to create an object of the class like,
`model = BayesianRegression(4,2)
It works fine for out_features = 1. I want to keep it bigger than one.
How do I modify the class so that it works?

Hi @mxnyy, does this work for you?

  class BayesianRegression(PyroModule):
      def __init__(self, in_features, out_features):
          super().__init__()
          self.linear = PyroModule[nn.Linear](in_features, out_features)
          self.linear.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2))
          self.linear.bias = PyroSample(dist.Normal(0., 1.).expand([out_features]).to_event(1))

      def forward(self, x, y=None):
-         sigma = pyro.sample("sigma", dist.Uniform(0., 10.))
+         sigma = pyro.sample("sigma", dist.Uniform(0., 10.)).unsqueeze(-1)
-         mean = self.linear(x).squeeze(-1)
+         mean = self.linear(x)
          with pyro.plate("data", x.shape[0]):
-             obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
+             obs = pyro.sample("obs", dist.Normal(mean, sigma).to_event(1), obs=y)
          return mean

Hi @fritzo, it does the work. Thanks a lot. I will run it for real data and see if it works for my task. Thanks again.

1 Like