<jupyter_start><jupyter_text><jupyter_code>import numpy as np
import torch
import torch.nn as nn 
import torch.nn.functional as F
import torch.optim as optim
from tqdm import tqdm
import matplotlib.pyplot as plt
from sklearn.feature_selection import mutual_info_regression
# data
def gen_x(num_samples):
    return np.sign(np.random.normal(0.,1.,[num_samples,1]))

def gen_y(x, var):
    return x + np.random.normal(0.,np.sqrt(var),[x.shape[0],1])
var = 0.2
num_samples = 1000000
x=gen_x(num_samples)
y=gen_y(x, var)
p_y_x=np.exp(-(y-x)**2/(2*var))
p_y_x_minus=np.exp(-(y+1)**2/(2*var))
p_y_x_plus=np.exp(-(y-1)**2/(2*var))
mi=np.average(np.log(p_y_x/(0.5*p_y_x_minus+0.5*p_y_x_plus)))
mi
hidden_size=10
n_epoch = 500
data_size = 20000

class MINE(nn.Module):
    def __init__(self, hidden_size=10):
        super(MINE, self).__init__()
        self.layers= nn.Sequential(nn.Linear(2, hidden_size),
                                   nn.ReLU(),
                                   nn.Linear(hidden_size,1))

    def forward(self, x, y):
        batch_size = x.size(0)
        tiled_x = torch.cat([x, x,], dim=0)
        idx = torch.randperm(batch_size)
        
        shuffled_y = y[idx]
        concat_y = torch.cat([y, shuffled_y], dim=0)
        
        inputs = torch.cat([tiled_x, concat_y], dim=1)
        logits = self.layers(inputs)

        pred_xy = logits[:batch_size]
        pred_x_y = logits[batch_size:]
        loss = -(torch.mean(pred_xy) 
               - torch.log(torch.mean(torch.exp(pred_x_y))))

        return loss

model = MINE(hidden_size)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
plot_loss = []
all_mi = []
for epoch in tqdm(range(n_epoch)):
    x_sample=gen_x(data_size)
    y_sample=gen_y(x_sample, var)
    
    x_sample = torch.from_numpy(x_sample).float()
    y_sample = torch.from_numpy(y_sample).float()
    loss = model(x_sample, y_sample)
    
    model.zero_grad()
    loss.backward()
    optimizer.step()
   
    all_mi.append(-loss.item())


fig, ax = plt.subplots()
ax.plot(range(len(all_mi)), all_mi, label='MINE estimate')
ax.plot([0, len(all_mi)], [mi,mi], label='True Mutual Information')
ax.set_xlabel('training steps')
ax.legend(loc='best')<jupyter_output>100%|██████████| 500/500 [00:05<00:00, 94.65it/s]
{'fraction_non_alphanumeric': 0.1590308370044053, 'fraction_numerical': 0.030837004405286344, 'mean_word_length': 3.2768361581920904, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '431475', 'license': 'Apache-2.0', 'n_tokens_mistral': 949, 'n_tokens_neox': 900, 'n_words': 147, 'path': 'mutual_information.ipynb', 'repo_name': 'seanie12/MINE', 'size': '2270'}