Smoothed InfoNCE¶
Setup¶
import torch
import torch.nn.functional as F
import os
import numpy as np
import matplotlib.pyplot as plt
import argparse
from src.utils import *
from src.mi_estimators import *
from torch.utils.tensorboard import SummaryWriter
SEED = 0
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
np.random.seed(SEED)
cuda = True if torch.cuda.is_available() else False
FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor
torch.set_default_tensor_type(FloatTensor)
parser = argparse.ArgumentParser()
parser.add_argument("--rho", type=float, default=0.9, help="coefficient of Gaussian")
parser.add_argument("--d", type=int, default=20, help="dimension of X & Y")
parser.add_argument("--sample_size", type=int, default=400, help="sample size")
parser.add_argument("--gamma", type=float, default=1e-10, help="clipping parameter")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--batch_size", type=int, default=40, help="size of the batches")
parser.add_argument("--lr", type=float, default=1e-3, help="adam: learning rate")
parser.add_argument("--hidden_dim", type=int, default=100, help="Hidden dimension")
parser.add_argument("--ma_rate", type=float, default=0.1, help="move average rate")
parser.add_argument("--ma_ef", type=float, default=1, help="move average ef")
parser.add_argument("--alpha", type=float, default=1e-6, help="smooth parameter")
parser.add_argument("--reg", type=int, default=1, help="if apply regularization")
parser.add_argument("--n_epoch", type=int, default=2000, help="number of epochs of training")
parser.add_argument("--n_iters_1epoch", type=int, default=10, help="number of epochs of training")
opt, unknown = parser.parse_known_args()
opt.n_iters = opt.n_epoch * opt.n_iters_1epoch
ma_rate = 0.01 # moving average rate
continue_train = False # set to True to continue to train
load_available = True # set to False to prevent loading previous results
overwrite = False # set to True to overwrite previously stored results
Generate the data.
data = GaussianData(opt.sample_size, d=opt.d, rho=opt.rho)
X, Y, XY, Ground_truth = data.X, data.Y, torch.cat((data.X, data.Y), dim=1), data.mutual_information()
# choose n marginal samples to train DT
ref_X, ref_Y = shuffle_data(X, Y, opt.sample_size)
ref_XY = torch.cat([ref_X, ref_Y], dim = 1)
train_data = torch.cat([XY, ref_XY], dim = 0)
train_label = torch.cat([torch.ones([opt.sample_size,1]), torch.zeros([opt.sample_size, 1])], dim=0)
Train Decision Tree
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
clf = DecisionTreeClassifier(min_samples_leaf=20, max_depth=20)
# clf = DecisionTreeClassifier(min_samples_leaf=400, max_depth=10)
# clf = RandomForestClassifier(n_estimators=10, min_samples_split=5)
clf.fit(train_data.cpu().numpy(), train_label.cpu().numpy())
DecisionTreeClassifier(max_depth=20, min_samples_leaf=20)
path = clf.cost_complexity_pruning_path(train_data.cpu().numpy(), train_label.cpu().numpy())
ccp_alphas, impurities = path.ccp_alphas, path.impurities
Choose the confidence factor
ccp = ccp_alphas[round(len(ccp_alphas)/2)]
clf = DecisionTreeClassifier(min_samples_leaf=20, max_depth=20, ccp_alpha=ccp)
clf.fit(train_data.cpu().numpy(), train_label.cpu().numpy())
DecisionTreeClassifier(ccp_alpha=0.0032447663535712136, max_depth=20,
min_samples_leaf=20)
The score of DT.
clf.score(train_data.cpu().numpy(), train_label.cpu().numpy())
0.7675
# construct the data needs to train and predict the label
x_tile = X.unsqueeze(0).repeat((opt.sample_size, 1, 1))
y_tile = Y.unsqueeze(1).repeat((1, opt.sample_size, 1))
data_matrix = torch.cat([x_tile, y_tile], dim = -1)
DT_prob_matrix = torch.Tensor(clf.predict_proba(data_matrix.reshape(-1, opt.d*2).cpu().numpy())[:,1].reshape(opt.sample_size, opt.sample_size))
Initialize the model.¶
name = f'./results/Smoothed_InfoNCE_dim{opt.d}_reg{opt.reg}_alpha{opt.alpha}' # filename
chkpt_name = name+'.pt' # checkpoint
from datetime import datetime
TIMESTAMP = "{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.now())
writer = SummaryWriter(f'./results/log/Smoothed_InfoNCE_dim{opt.d}_reg{opt.reg}_alpha{opt.alpha}/{TIMESTAMP}')
discriminator = Net(input_size=opt.d*2, hidden_layers=2, hidden_size=100)
# move NN model to GPU if GPU is available
if cuda:
discriminator.cuda()
# Adam optimizer
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
if load_available and os.path.exists(chkpt_name):
checkpoint = torch.load(
chkpt_name, map_location='cuda' if torch.cuda.is_available() else 'cpu')
mi_list = checkpoint['mi_list']
model_state = checkpoint['model_state']
discriminator.load_state_dict(model_state)
print('Previous results loaded.')
else:
mi_list = [] # storing the mi estimation of each iteration
Previous results loaded.
def smooth_loss(net, prob_matrix, x_data, y_data, alpha, writer=None, epoch=None, reg=False):
# alpha: smoothing factor
# net: torch.nn.Module
# prob_matrix: the probability estimate matrix of given batch data. size: batchsize * batchsize
# return: loss of smoothed infoNCE
# first term:
# E[f(X,Y)*p(C=1|X,Y)] + E[f(X',Y')*p(C=1|X',Y')]
joint_samples = torch.cat((x_data, y_data), dim=1)
x_tile = x_data.unsqueeze(0).repeat((x_data.shape[0], 1, 1))
y_tile = y_data.unsqueeze(1).repeat((1, x_data.shape[0], 1))
train_data = torch.cat([x_tile, y_tile], dim = -1).reshape(-1, x_data.shape[1]+y_data.shape[1])
train_label = torch.eye(x_data.shape[0]).reshape(-1,)
# construct the marginal samples
marginal_samples = train_data[train_label==0,:]
prob_DT = prob_matrix.reshape(-1, 1)
# smooth the hard labels with the probability of DT
hard_label = train_label.reshape(-1,1)
prob_DT = alpha*prob_DT + hard_label*(1-alpha)
pos_prob_DT = prob_DT[train_label==1,:]
neg_prob_DT = prob_DT[train_label==0,:]
a = (net(joint_samples)*pos_prob_DT).mean() + (net(marginal_samples)*neg_prob_DT).mean()
# second term:
# E[log (E[e^f(X,Y')*p(C=0|X,Y')*P(C=0)|X] + E[e^f(X,Y)*p(C=0|X,Y)*p(C=1)|X]) - log p(C=0|X) ]
b_list = []
mask = torch.eye(x_data.shape[0])
for i in range(x_data.shape[0]):
x_i = x_data[i,:]
batch_label = torch.zeros(x_data.shape[0])
batch_label[i] = 1
x_tile = x_i.unsqueeze(0).repeat((x_data.shape[0], 1))
batch_xy = torch.cat([x_tile, y_data], dim = 1)
# constuct X_tilde Y_tilde by repeating (x_i,y_i) n-1 times and concatenate with all cross samples,
# therefore there are 2*(n-1) samples in total for a given x_i
batch_xy_ = torch.cat((batch_xy[batch_label==1,:].repeat((batch_xy.shape[0]-1, 1)), batch_xy[mask[i]==0,:]), dim=0)
# P(C=0|x) = E[P(C=0|x,Y')]
# pcx is the estimate of p(C=0|x)
prob_ = torch.cat((prob_matrix[i, i].repeat(batch_xy.shape[0]-1), prob_matrix[i,mask[i]==0]), dim=0)
# hard_label_ here is the hard label of probability that p(C=0|X,Y)
hard_label_ = torch.cat((torch.zeros(x_data.shape[0]-1), torch.ones(x_data.shape[0]-1))).reshape(-1,1)
# pcx = (alpha*prob_.reshape(-1,1) + hard_label_*(1-alpha)).mean()
prob_xy_ = alpha*prob_.reshape(-1,1) + hard_label_*(1-alpha)
# b_list.append(torch.logsumexp(net(batch_xy_)+torch.log(prob_xy_), dim=0) - np.log(batch_xy_.shape[0]) - torch.log(pcx))
b_list.append(torch.logsumexp(net(batch_xy_)+torch.log(prob_xy_), dim=0) - np.log(batch_xy_.shape[0]))
if writer is not None:
writer.add_scalar('a', a, epoch)
writer.add_scalar('b', torch.mean(torch.stack(b_list)), epoch)
if reg:
return a - sum(b_list)/len(b_list) - (sum(b_list)/len(b_list))**2
else:
return a - sum(b_list)/len(b_list)
if continue_train:
_iter = 0
for i in range(opt.n_epoch):
idx = np.random.permutation(opt.sample_size)
for j in range(opt.n_iters_1epoch):
batch_idx = idx[j::opt.n_iters_1epoch]
batch_X = X[batch_idx]
batch_Y = Y[batch_idx]
prob_batch = DT_prob_matrix[np.ix_(batch_idx, batch_idx)]
optimizer_D.zero_grad()
loss = - smooth_loss(discriminator, prob_batch, batch_X, batch_Y, alpha=opt.alpha, reg=opt.reg) # negative infonce_bound as the loss
loss.backward()
optimizer_D.step()
with torch.no_grad():
# mi_est = smooth_loss(discriminator, clf, X, Y, alpha=0.2)
mi_est = smooth_loss(discriminator, DT_prob_matrix, X, Y, alpha=opt.alpha)
mi_list.append(mi_est.item())
writer.add_scalar('mi_list', mi_est.item(), _iter)
writer.add_scalar('loss', loss, _iter)
_iter += 1
if _iter%200==0:
print("Iternation: %d, loss: %f, mi_est: %f"%(_iter, loss.item(), mi_est))
fig = plot_fig(discriminator, X, Y, opt.d if opt.d<6 else 6)
writer.add_figure('heatmap', fig, _iter)
writer.add_histogram('first layer', discriminator.fc[0].weight.data, _iter)
writer.add_histogram('second layer', discriminator.fc[1].weight.data, _iter)
writer.add_histogram('third layer', discriminator.fc[2].weight.data, _iter)
writer.add_histogram('first layer (grad)', discriminator.fc[0].weight.grad.data, _iter)
writer.add_histogram('second layer (grad)', discriminator.fc[1].weight.grad.data, _iter)
writer.add_histogram('third layer (grad)', discriminator.fc[2].weight.grad.data, _iter)
writer.add_graph(discriminator, (XY,))
writer.close()
if overwrite or not os.path.exists(chkpt_name):
model_state = discriminator.state_dict()
torch.save({
'mi_list': mi_list,
'model_state': model_state
}, chkpt_name)
writer.close()
print('Current results saved.')
mi_copy = mi_list.copy()
for k in range(1,len(mi_list)):
mi_copy[k] = (1-ma_rate) * mi_copy[k-1] + ma_rate * mi_copy[k]
KSG estimator
import src.entropy_estimators as ee
KSG = ee.mi(X.cpu().numpy(),Y.cpu().numpy())
plt.plot(mi_copy, label='Smoothed InfoNCE')
plt.axhline(Ground_truth,label='Ground Truth',linestyle='--',color='red')
plt.axhline(KSG,label='KSG',linestyle='--',color='blue')
plt.axhline(np.log(opt.sample_size),label='log $n$',linestyle='--',color='grey')
plt.xlabel('number of iterations')
plt.ylabel('MI estimation')
# plt.title('MI estimation with Smoothed InfoNCE')
plt.legend()
# plt.savefig(f'results/Smoothed_InfoNCE_dim{opt.d}_ma{ma_rate}.pdf')
<matplotlib.legend.Legend at 0x7f257df0cdd0>
heatmap = plot_fig(discriminator, X, Y, vmin=-80, vmax=20, d=6)
# plt.savefig(f'./results/heatmap_smoothed_InfoNCE_dim{opt.d}.pdf')
