[Pytorch neural network practical case] 38 neural network realizes the function of estimating mutual information

1 Case description (implementing the function of MINE positive method)

Define two sets of simulated data with different distributions, and use the MINE method of neural networks to calculate the mutual information between the two data distributions

2 Code writing

2.1 Code combat: prepare sample data

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'  # 可能是由于是MacOS系统的原因

### 本例实现了用神经网络计算互信息的功能。这是一个简单的例子，目的在于帮助读者更好地理MNE方法。

# 1.1 准备样本数据：定义两个数据生成函数gen_x()、gen_y()。函数gen_x()用于生成1或-1，函数gen_y()在此基础上为其再加上一个符合高斯分布的随机值。
# 生成模拟数据
data_size = 1000
def gen_x():
    return np.sign(np.random.normal(0.0,1.0,[data_size,1]))

def gen_y(x):
    return x + np.random.normal(0.0,0.5,[data_size,1])

def show_data():
    x_sample = gen_x()
    y_sample = gen_y(x_sample)
    plt.scatter(np.arange(len(x_sample)), x_sample, s=10,c='b',marker='o')
    plt.scatter(np.arange(len(y_sample)), y_sample, s=10,c='y',marker='o')
    plt.show() # 两条横线部分是样本数据x中的点，其他部分是样本数据y。

2.2 Code combat: Defining the neural network model

# 1.2 定义神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(1,10)
        self.fc2 = nn.Linear(1,10)
        self.fc3 = nn.Linear(10,1)

    def forward(self,x,y):
        h1 = F.relu(self.fc1(x) + self.fc2(y))
        h2 = self.fc3(h1)
        return h2

2.3 Code combat: use the MINE method to train the model and output the results

# 1.3 利用MINE方法训练模型并输出结果
if __name__ == '__main__':
    show_data()# 显示数据
    model = Net() # 实例化模型
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01) # 使用Adam优化器并设置学习率为0.01

    n_epoch = 500
    plot_loss = []
    # MiNE方法主要用于模型的训练阶段
    for epoch in tqdm(range(n_epoch)):
        x_sample = gen_x() # 调用gen_x()函数生成样本x_Sample。X_sample代表X的边缘分布P(X)
        y_sample = gen_y(x_sample) # 将生成的×_sample样本放到gen_x()函数中，生成样本y_sample。y_sample代表条件分布P(Y|X)。
        y_shuffle = np.random.permutation(y_sample) # )将 y_sample按照批次维度打乱顺序得到y_shuffle，y_shuffle是Y的经验分布，近似于Y的边缘分布P(Y)。
        # 转化为张量
        x_sample = torch.from_numpy(x_sample).type(torch.FloatTensor)
        y_sample = torch.from_numpy(y_sample).type(torch.FloatTensor)
        y_shuffle = torch.from_numpy(y_shuffle).type(torch.FloatTensor)

        model.zero_grad()
        pred_xy = model(x_sample, y_sample)  # 式(8-49）中的第一项联合分布的期望:将x_sample和y_sample放到模型中，得到联合概率（P(X,Y)=P(Y|X)P(X)）关于神经网络的期望值pred_xy。
        pred_x_y = model(x_sample, y_shuffle)  # 式(8-49)中的第二项边缘分布的期望:将x_sample和y_shuffle放到模型中，得到边缘概率关于神经网络的期望值pred_x_y 。

        ret = torch.mean(pred_xy) - torch.log(torch.mean(torch.exp(pred_x_y))) # 将pred_xy和pred_x_y代入式（8-49）中，得到互信息ret。
        loss = - ret  # 最大化互信息：在训练过程中，因为需要将模型权重向着互信息最大的方向优化，所以对互信息取反，得到最终的loss值。
        plot_loss.append(loss.data)  # 收集损失值
        loss.backward()  # 反向传播：在得到loss值之后，便可以进行反向传播并调用优化器进行模型优化。
        optimizer.step()  # 调用优化器
    plot_y = np.array(plot_loss).reshape(-1, )  # 可视化
    plt.plot(np.arange(len(plot_loss)), -plot_y, 'r') # 直接将|oss值取反，得到最大化互信息的值。
    plt.show()

3 Code overview

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'  # 可能是由于是MacOS系统的原因

### 本例实现了用神经网络计算互信息的功能。这是一个简单的例子，目的在于帮助读者更好地理MNE方法。

# 1.1 准备样本数据：定义两个数据生成函数gen_x()、gen_y()。函数gen_x()用于生成1或-1，函数gen_y()在此基础上为其再加上一个符合高斯分布的随机值。
# 生成模拟数据
data_size = 1000
def gen_x():
    return np.sign(np.random.normal(0.0,1.0,[data_size,1]))

def gen_y(x):
    return x + np.random.normal(0.0,0.5,[data_size,1])

def show_data():
    x_sample = gen_x()
    y_sample = gen_y(x_sample)
    plt.scatter(np.arange(len(x_sample)), x_sample, s=10,c='b',marker='o')
    plt.scatter(np.arange(len(y_sample)), y_sample, s=10,c='y',marker='o')
    plt.show() # 两条横线部分是样本数据x中的点，其他部分是样本数据y。

# 1.2 定义神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(1,10)
        self.fc2 = nn.Linear(1,10)
        self.fc3 = nn.Linear(10,1)

    def forward(self,x,y):
        h1 = F.relu(self.fc1(x) + self.fc2(y))
        h2 = self.fc3(h1)
        return h2

# 1.3 利用MINE方法训练模型并输出结果
if __name__ == '__main__':
    show_data()# 显示数据
    model = Net() # 实例化模型
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01) # 使用Adam优化器并设置学习率为0.01

    n_epoch = 500
    plot_loss = []
    # MiNE方法主要用于模型的训练阶段
    for epoch in tqdm(range(n_epoch)):
        x_sample = gen_x() # 调用gen_x()函数生成样本x_Sample。X_sample代表X的边缘分布P(X)
        y_sample = gen_y(x_sample) # 将生成的×_sample样本放到gen_x()函数中，生成样本y_sample。y_sample代表条件分布P(Y|X)。
        y_shuffle = np.random.permutation(y_sample) # )将 y_sample按照批次维度打乱顺序得到y_shuffle，y_shuffle是Y的经验分布，近似于Y的边缘分布P(Y)。
        # 转化为张量
        x_sample = torch.from_numpy(x_sample).type(torch.FloatTensor)
        y_sample = torch.from_numpy(y_sample).type(torch.FloatTensor)
        y_shuffle = torch.from_numpy(y_shuffle).type(torch.FloatTensor)

        model.zero_grad()
        pred_xy = model(x_sample, y_sample)  # 式(8-49）中的第一项联合分布的期望:将x_sample和y_sample放到模型中，得到联合概率（P(X,Y)=P(Y|X)P(X)）关于神经网络的期望值pred_xy。
        pred_x_y = model(x_sample, y_shuffle)  # 式(8-49)中的第二项边缘分布的期望:将x_sample和y_shuffle放到模型中，得到边缘概率关于神经网络的期望值pred_x_y 。

        ret = torch.mean(pred_xy) - torch.log(torch.mean(torch.exp(pred_x_y))) # 将pred_xy和pred_x_y代入式（8-49）中，得到互信息ret。
        loss = - ret  # 最大化互信息：在训练过程中，因为需要将模型权重向着互信息最大的方向优化，所以对互信息取反，得到最终的loss值。
        plot_loss.append(loss.data)  # 收集损失值
        loss.backward()  # 反向传播：在得到loss值之后，便可以进行反向传播并调用优化器进行模型优化。
        optimizer.step()  # 调用优化器
    plot_y = np.array(plot_loss).reshape(-1, )  # 可视化
    plt.plot(np.arange(len(plot_loss)), -plot_y, 'r') # 直接将|oss值取反，得到最大化互信息的值。
    plt.show()