# Fit Gaussian Mixture Model

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()
``````

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.

Thanks a lot !

Hi Luc,

The scale, so std, shouldn’t be negative nor very large.

One way to circumvent this is to use:

``````std_weight = 1e-4
comp = D.Normal(self.means, std_weight * self.stdevs.abs())
``````

And for start: I would update the stdevs at all. In that way, you only train the means and weights, which in some applications works.

Best,
Tessa

1 Like