Fit Gaussian Mixture Model

1

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 !

1 Like
2

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