tsne_nn.py

import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
from openTSNE import TSNE
import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import SpectralEmbedding
from scipy.sparse import save_npz, load_npz
from functools import partial

EPS = 1e-12
BATCH_SIZE = 256
N_EPOCHS = 200
DATA_NPZ_PATH = 'mnist_60000.npz'

def get_activation(act):
    if act == 'lrelu':
        return nn.LeakyReLU(0.2, inplace=True)
    elif act == 'relu':
        return nn.ReLU(inplace=True)
    raise Exception('unsupported activation function')
    
class FCEncoder(nn.Module):
    def __init__(self, dim, num_layers=3, act='lrelu'):
        super(FCEncoder, self).__init__()
        self.dim = dim
        self.num_layers = num_layers
        self.act = partial(get_activation, act=act)
        hidden_dim = 256
        layers = [
            (nn.Linear(dim, hidden_dim*2)),
            self.act(),
            (nn.Linear(hidden_dim*2, hidden_dim)),
            self.act(),
        ]
        layers += [
            (nn.Linear(hidden_dim, hidden_dim)),
            self.act(),
        ] * num_layers
        layers += [
            (nn.Linear(hidden_dim, 2)),
        ]
        self.net = nn.Sequential(*layers)
        
    def forward(self, X):
        return self.net(X)

def neg_squared_euc_dists(X):
    sum_X = X.pow(2).sum(dim=1)
    D = (-2 * X @ X.transpose(1, 0) + sum_X).transpose(1, 0) + sum_X
    return -D

def w_tsne(Y):
    distances = neg_squared_euc_dists(Y)
    inv_distances = (1. - distances).pow(-1) #1 / (1+d^2)
    return inv_distances

def KLD(P, Q):
    return P * torch.log((P+EPS) / Q)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device:', device)

print('load data')
mnist = np.load(DATA_NPZ_PATH)
data = mnist['data']

print('calc P')
try:
    P_csc = load_npz('P_tsne_csc.npz')
    print('Use P cache')
except:
    print('Use new P')
    pre_embedding = TSNE(perplexity=30).prepare_initial(data)
    P_csc = pre_embedding.affinities.P
    save_npz('P_tsne_csc', pre_embedding.affinities.P)


print('Trying to put X into GPU')
X = torch.from_numpy(data).float()
X = X.to(device)

print('Constructing NN')
encoder = FCEncoder(784, num_layers=3)
encoder = encoder.to(device)
encoder = encoder.float()

init_lr = 1e-3
optimizer = optim.Adam(encoder.parameters(), lr=init_lr, betas=(0, 0.9))
# optimizer = optim.SGD(encoder.parameters(), lr=init_lr)

print('optimizing...')
for epoch in range(N_EPOCHS):
    idxs = torch.randperm(len(X))
    for i in range(0, len(X), BATCH_SIZE):
        idx = idxs[i:i+BATCH_SIZE]
        p = torch.Tensor(P_csc[idx][:, idx].toarray()).float().to(device)
        optimizer.zero_grad()
        y = encoder(X[idx])
        w = w_tsne(y)
        q = w / torch.sum(w)
        loss = KLD(p, q).sum()
        loss.backward()
        optimizer.step()
        
    print('\r', '{:03d}'.format(epoch), '{:.5f}'.format(loss.item()), end='\n')
    with torch.no_grad():
        Y = encoder(X)
    new_lr = (1 - epoch / N_EPOCHS) * init_lr
    for param_group in optimizer.param_groups:
        param_group['lr'] = new_lr
#     print('Y saved.', 'KLD(P, Q)')
    np.savez_compressed('tsne_Y', Y=Y.detach().cpu().numpy())