C++高斯混合型GMM实现

#ifndef MYGMM_H
#define MYGMM_H
#include <iostream>
#include <Eigen/Dense>
#include <math.h>
#include <vector>
#define EPS 0.00000001

using namespace std;
using namespace Eigen;

Eigen::MatrixXf getCov(Eigen::MatrixXf input)
{
    // 按列求均值
    Eigen::MatrixXf meanVec = input.colwise().mean();
    Eigen::RowVectorXf meanVecRow(Eigen::RowVectorXf::Map(meanVec.data(), input.cols()));
    Eigen::MatrixXf zeroMeanMat = input;
    zeroMeanMat.rowwise() -= meanVecRow;
    if (input.rows() == 1)
        return (zeroMeanMat.adjoint() * zeroMeanMat) / double(input.rows());
    else
        return (zeroMeanMat.adjoint() * zeroMeanMat) / double(input.rows() - 1);
}

float GMM(Eigen::RowVectorXf u, Eigen::MatrixXf sigma, Eigen::RowVectorXf x)
{
    float tmp = -0.5 * (x - u) * sigma.inverse() * (x - u).transpose();
    return exp(tmp) / pow(2.0 * 3.1416, sigma.rows() / 2.0) / pow(sigma.determinant(), 0.5);
}

class MyGMM
{
public:
    void Init(vector<vector<float>> input)
    {
        Eigen::MatrixXf tmp(input.size(), input[0].size());
        input_ = tmp;
        rows_ = input_.rows();
        cols_ = input_.cols();
        for (int i = 0; i < rows_; ++i)
        {
            for (int j = 0; j < cols_; ++j)
            {
                input_(i, j) = input[i][j];
            }
        }
        K_ = 3;
        a_.resize(K_);
        float kk = 0;
        for (int i = 0; i < K_ - 1; ++i)
        {
            a_[i] = 1.0 / K_;
            kk += a_[i];
        }
        a_[K_ - 1] = 1 - kk;
        Eigen::MatrixXf u(K_, cols_);
        sigma_.resize(K_);
        for (int i = 0; i < K_; ++i)
        {
            u.row(i) = input_.row(rows_ * i / K_);
            sigma_[i] = getCov(input_);
        }
        u_ = u;
        belong_.resize(rows_);
        iter_ = 0;
        last_ = 0;
    }

    int EM()
    {
        while (1)
        {
            vector<int> N(K_, 0);
            Eigen::RowVectorXf gamma = RowVectorXf::Zero(rows_);
            // Expectation
            for (int i = 0; i < rows_; ++i)
            {
                float sum = 0;
                for (int j = 0; j < K_; ++j)
                {
                    if (a_[j] < EPS)
                        continue;
                    float tmp = a_[j] * GMM(u_.row(j), sigma_[j], input_.row(i));
                    sum += tmp;
                    if (tmp > gamma[i])
                    {
                        gamma[i] = tmp;
                        belong_[i] = j;
                    }
                }
                gamma[i] /= sum;
                N[belong_[i]]++;
            }
            // 
            for (int j = 0; j < K_; ++j)
            {
                if (N[j] == 0)
                {
                    return 0;
                } else if(N[j] == 1) {
                    return 1;
                }
            }
            // reset
            u_ = MatrixXf::Zero(K_, cols_);
            for (int j = 0; j < K_; ++j)
            {
                sigma_[j] = MatrixXf::Zero(cols_, cols_);
            }
            // Maximization
            for (int i = 0; i < rows_; ++i)
            {
                u_.row(belong_[i]) += gamma[i] * input_.row(i);
            }
            for (int j = 0; j < K_; ++j)
            {
                u_.row(j) /= N[j];
            }
            for (int i = 0; i < rows_; ++i)
            {
                sigma_[belong_[i]] += gamma[i] * (input_.row(i) - u_.row(belong_[i])).transpose() * (input_.row(i) - u_.row(belong_[i]));
            }
            for (int j = 0; j < K_; ++j)
            {
                sigma_[j] /= N[j];
                if (sigma_[j].determinant() == 0)
                {
                    return 2;
                }
                a_[j] = N[j] * 1.0 / rows_;
            }
            // Judge
            float judge = 0;
            for (int i = 0; i < rows_; ++i)
            {
                int tmp = belong_[i];
                float gg = GMM(u_.row(tmp), sigma_[tmp], input_.row(i));
                judge = judge + log(a_[tmp]) + log(gg);
            }
            ++iter_;
            // cout << judge << endl;
            if (iter_ > 6 && fabs(last_ - judge) < 0.00001)
                return 9;
            last_ = judge;
        }
    }

    Eigen::MatrixXf input_;
    int rows_;
    int cols_;
    int K_;
    vector<float> a_;
    Eigen::MatrixXf u_;
    vector<MatrixXf> sigma_;
    vector<int> belong_;
    int iter_;
    float last_;
};
#endif