EM算法解决二硬币问题(python)

前言

EM算法是很多深度学习的基础，这里我们以一个简单的例子去探寻EM算法的原理。

关于EM算法的原理，可见: EM算法详解 (知乎)

不断更新，先挖坑
代码来源: https://github.com/luwill/Machine_Learning_Code_Implementation/tree/master/charpter22_EM

问题描述

有两个材质不均匀的硬币(正反面概率不是0.5), 分别是硬币B和硬币C，随机选取硬币， 进行5组独立实验，每组实验进行10次，每次随机选取一个硬币进行抛掷，记录下正反面次数分别为 (5,5), (9,1), (8,2), (4,6), (7,3)。
你能估计最可能的硬币B和硬币C抛掷正反面概率分别为多少？

源码

EM算法先假设硬币B, 硬币C的概率为一个初始化值(可以随机，也可以由我们看数据分布手动设定)， 例如这里我们设置分别为0.6和0.5。

# 导入numpy库 
import numpy as np

### EM算法过程函数定义
def em(data, thetas, max_iter=30, eps=1e-3):
    '''
    输入：
    data：观测数据
    thetas：初始化的估计参数值
    max_iter：最大迭代次数
    eps：收敛阈值
    输出：
    thetas：估计参数
    '''
    # 初始化似然函数值
    ll_old = -np.infty
    for i in range(max_iter):
        ### E步：求隐变量分布
        # 对数似然 [coin_num, exp_num], [2, 5]
        log_like = np.array([np.sum(data * np.log(theta), axis=1) for theta in thetas])
        # 似然 [coin_num, exp_num], [2, 5]
        like = np.exp(log_like)
        # 求隐变量分布 [coin_num, exp_num], [2, 5]
        ws = like/like.sum(0)
        # 概率加权 [2, 5, 2]
        vs = np.array([w[:, None] * data for w in ws])
        ### M步：更新参数值 [2, 2]
        thetas = np.array([v.sum(0)/v.sum() for v in vs])
        # 更新似然函数 
        ll_new = np.sum([w*l for w, l in zip(ws, log_like)])
        print("Iteration: %d" % (i+1))
        print("theta_B = %.2f, theta_C = %.2f, ll = %.2f" 
              % (thetas[0,0], thetas[1,0], ll_new))
        # 满足迭代条件即退出迭代
        if np.abs(ll_new - ll_old) < eps:
            break
        ll_old = ll_new
    return thetas


if __name__ == "__main__":
    # 观测数据，5次独立试验，每次试验10次抛掷的正反次数
    # 比如第一次试验为5次正面5次反面
    observed_data = np.array([(5,5), (9,1), (8,2), (4,6), (7,3)])
    # 初始化参数值，即硬币B的正面概率为0.6，硬币C的正面概率为0.5
    thetas = np.array([[0.6, 0.4], [0.5, 0.5]])
    # thetas = em(observed_data, thetas, max_iter=30, eps=1e-3)
    thetas = em(observed_data, thetas, max_iter=30, eps=1e-4)
    print(thetas)

由结果我们发现， 硬币B正面概率约为为0.7968， 硬币C正面的概率约为0.5196。