Hello !
I’d like to train a very basic Mixture of 2 Gaussians to segment background in a 2d image.
However I think I’m confused on how to use torch.distribution.
This is what I’m doing:

first I prepare my 2d numpy array by doing:

x = torch.from_numpy(image.reshape((image.size, 1)))


then I define a Module as bellow:

class GaussianMixtureModel(torch.nn.Module):

    def __init__(self, n_components: int=2):
        super().__init__()
        weights = torch.ones(n_components, )
        means   = torch.randn(n_components, )
        stdevs  = torch.tensor(np.abs(np.random.randn(n_components, )))
        self.weights = torch.nn.Parameter(weights)
        self.means   = torch.nn.Parameter(means)
        self.stdevs  = torch.nn.Parameter(stdevs)
    
    def forward(self, x):
        mix  = D.Categorical(self.weights)
        comp = D.Normal(self.means, self.stdevs)
        gmm  = D.MixtureSameFamily(mix, comp)
        return - gmm.log_prob(x).mean()


then I define a basic training procedure:

model = GaussianMixtureModel(n_components=2)
optim = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

n_iter = 1_000
for _ in range(n_iter):
    loss = model(x)
    loss.backward()
    optim.step()
    optim.zero_grad()

But I get either:

The parameter probs has invalid values (error caught when defining mix in the forward method)
The parameters scale has invalid values (error caught when defining comp in the forward method

So I guess I’m doing something wrong with my distributions definitions.
Could you please help me ?
Thanks a lot !