263 lines
7.5 KiB
Python
263 lines
7.5 KiB
Python
#
|
|
# tsne_torch.py
|
|
#
|
|
# Implementation of t-SNE in pytorch. The implementation was tested on pytorch
|
|
# > 1.0, and it requires Numpy to read files. In order to plot the results,
|
|
# a working installation of matplotlib is required.
|
|
#
|
|
#
|
|
# The example can be run by executing: `python tsne_torch.py`
|
|
#
|
|
#
|
|
# Created by Xiao Li on 23-03-2020.
|
|
# Copyright (c) 2020. All rights reserved.
|
|
from random import shuffle
|
|
|
|
import numpy as np
|
|
import matplotlib.pyplot as pyplot
|
|
import argparse
|
|
import torch
|
|
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
|
|
from tqdm import tqdm
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--xfile", type=str, default="mnist2500_X.txt", help="file name of feature stored")
|
|
parser.add_argument("--cuda", type=int, default=1, help="if use cuda accelarate")
|
|
|
|
opt = parser.parse_args()
|
|
print("get choice from args", opt)
|
|
xfile = opt.xfile
|
|
|
|
if opt.cuda:
|
|
print("set use cuda")
|
|
torch.set_default_tensor_type(torch.cuda.DoubleTensor)
|
|
else:
|
|
torch.set_default_tensor_type(torch.DoubleTensor)
|
|
|
|
|
|
def Hbeta_torch(D, beta=1.0):
|
|
P = torch.exp(-D.clone() * beta)
|
|
|
|
sumP = torch.sum(P)
|
|
|
|
H = torch.log(sumP) + beta * torch.sum(D * P) / sumP
|
|
P = P / sumP
|
|
|
|
return H, P
|
|
|
|
|
|
def x2p_torch(X, tol=1e-5, perplexity=30.0):
|
|
"""
|
|
Performs a binary search to get P-values in such a way that each
|
|
conditional Gaussian has the same perplexity.
|
|
"""
|
|
|
|
# Initialize some variables
|
|
print("Computing pairwise distances...")
|
|
(n, d) = X.shape
|
|
|
|
sum_X = torch.sum(X*X, 1)
|
|
D = torch.add(torch.add(-2 * torch.mm(X, X.t()), sum_X).t(), sum_X)
|
|
|
|
P = torch.zeros(n, n)
|
|
beta = torch.ones(n, 1)
|
|
logU = torch.log(torch.tensor([perplexity]))
|
|
n_list = [i for i in range(n)]
|
|
|
|
# Loop over all datapoints
|
|
for i in range(n):
|
|
|
|
# Print progress
|
|
if i % 500 == 0:
|
|
print("Computing P-values for point %d of %d..." % (i, n))
|
|
|
|
# Compute the Gaussian kernel and entropy for the current precision
|
|
# there may be something wrong with this setting None
|
|
betamin = None
|
|
betamax = None
|
|
Di = D[i, n_list[0:i]+n_list[i+1:n]]
|
|
|
|
(H, thisP) = Hbeta_torch(Di, beta[i])
|
|
|
|
# Evaluate whether the perplexity is within tolerance
|
|
Hdiff = H - logU
|
|
tries = 0
|
|
while torch.abs(Hdiff) > tol and tries < 50:
|
|
|
|
# If not, increase or decrease precision
|
|
if Hdiff > 0:
|
|
betamin = beta[i].clone()
|
|
if betamax is None:
|
|
beta[i] = beta[i] * 2.
|
|
else:
|
|
beta[i] = (beta[i] + betamax) / 2.
|
|
else:
|
|
betamax = beta[i].clone()
|
|
if betamin is None:
|
|
beta[i] = beta[i] / 2.
|
|
else:
|
|
beta[i] = (beta[i] + betamin) / 2.
|
|
|
|
# Recompute the values
|
|
(H, thisP) = Hbeta_torch(Di, beta[i])
|
|
|
|
Hdiff = H - logU
|
|
tries += 1
|
|
|
|
# Set the final row of P
|
|
P[i, n_list[0:i]+n_list[i+1:n]] = thisP
|
|
|
|
# Return final P-matrix
|
|
return P
|
|
|
|
|
|
def pca_torch(X, no_dims=50):
|
|
print("Preprocessing the data using PCA...")
|
|
(n, d) = X.shape
|
|
X = X - torch.mean(X, 0)
|
|
|
|
(l, M) = torch.eig(torch.mm(X.t(), X), True)
|
|
# split M real
|
|
for i in range(d):
|
|
if l[i, 1] != 0:
|
|
M[:, i+1] = M[:, i]
|
|
i += 1
|
|
|
|
Y = torch.mm(X, M[:, 0:no_dims])
|
|
return Y
|
|
|
|
|
|
def tsne(X, no_dims=2, initial_dims=50, perplexity=30.0):
|
|
"""
|
|
Runs t-SNE on the dataset in the NxD array X to reduce its
|
|
dimensionality to no_dims dimensions. The syntaxis of the function is
|
|
`Y = tsne.tsne(X, no_dims, perplexity), where X is an NxD NumPy array.
|
|
"""
|
|
|
|
# Check inputs
|
|
if isinstance(no_dims, float):
|
|
print("Error: array X should not have type float.")
|
|
return -1
|
|
if round(no_dims) != no_dims:
|
|
print("Error: number of dimensions should be an integer.")
|
|
return -1
|
|
|
|
# Initialize variables
|
|
X = pca_torch(X, initial_dims).to('cuda') # Sending to('cuda') after because torch.eig is broken in Windows currently on Ampere GPUs.
|
|
(n, d) = X.shape
|
|
max_iter = 1000
|
|
initial_momentum = 0.5
|
|
final_momentum = 0.8
|
|
eta = 500
|
|
min_gain = 0.01
|
|
Y = torch.randn(n, no_dims)
|
|
dY = torch.zeros(n, no_dims)
|
|
iY = torch.zeros(n, no_dims)
|
|
gains = torch.ones(n, no_dims)
|
|
|
|
# Compute P-values
|
|
P = x2p_torch(X, 1e-5, perplexity)
|
|
P = P + P.t()
|
|
P = P / torch.sum(P)
|
|
P = P * 4. # early exaggeration
|
|
print("get P shape", P.shape)
|
|
P = torch.max(P, torch.tensor([1e-21]))
|
|
|
|
# Run iterations
|
|
for iter in tqdm(range(max_iter)):
|
|
|
|
# Compute pairwise affinities
|
|
sum_Y = torch.sum(Y*Y, 1)
|
|
num = -2. * torch.mm(Y, Y.t())
|
|
num = 1. / (1. + torch.add(torch.add(num, sum_Y).t(), sum_Y))
|
|
num[range(n), range(n)] = 0.
|
|
Q = num / torch.sum(num)
|
|
Q = torch.max(Q, torch.tensor([1e-12]))
|
|
|
|
# Compute gradient
|
|
PQ = P - Q
|
|
for i in range(n):
|
|
dY[i, :] = torch.sum((PQ[:, i] * num[:, i]).repeat(no_dims, 1).t() * (Y[i, :] - Y), 0)
|
|
|
|
# Perform the update
|
|
if iter < 20:
|
|
momentum = initial_momentum
|
|
else:
|
|
momentum = final_momentum
|
|
|
|
gains = (gains + 0.2) * ((dY > 0.) != (iY > 0.)).double() + (gains * 0.8) * ((dY > 0.) == (iY > 0.)).double()
|
|
gains[gains < min_gain] = min_gain
|
|
iY = momentum * iY - eta * (gains * dY)
|
|
Y = Y + iY
|
|
Y = Y - torch.mean(Y, 0)
|
|
|
|
# Compute current value of cost function
|
|
if (iter + 1) % 10 == 0:
|
|
C = torch.sum(P * torch.log(P / Q))
|
|
print("Iteration %d: error is %f" % (iter + 1, C))
|
|
|
|
# Stop lying about P-values
|
|
if iter == 100:
|
|
P = P / 4.
|
|
|
|
# Return solution
|
|
return Y
|
|
|
|
|
|
def run_tsne():
|
|
print("Run Y = tsne.tsne(X, no_dims, perplexity) to perform t-SNE on your dataset.")
|
|
|
|
limit = 4000
|
|
X, files = torch.load('results.pth')
|
|
zipped = list(zip(X, files))
|
|
shuffle(zipped)
|
|
X, files = zip(*zipped)
|
|
X = torch.cat(X, dim=0).squeeze()[:limit]
|
|
labels = np.zeros(X.shape[0]) # We don't have any labels..
|
|
|
|
# confirm that x file get same number point than label file
|
|
# otherwise may cause error in scatter
|
|
assert(len(X[:, 0])==len(X[:,1]))
|
|
assert(len(X)==len(labels))
|
|
|
|
with torch.no_grad():
|
|
Y = tsne(X, 2, 2048, 20.0)
|
|
|
|
if opt.cuda:
|
|
Y = Y.cpu().numpy()
|
|
|
|
# You may write result in two files
|
|
# print("Save Y values in file")
|
|
# Y1 = open("y1.txt", 'w')
|
|
# Y2 = open('y2.txt', 'w')
|
|
# for i in range(Y.shape[0]):
|
|
# Y1.write(str(Y[i,0])+"\n")
|
|
# Y2.write(str(Y[i,1])+"\n")
|
|
|
|
pyplot.scatter(Y[:, 0], Y[:, 1], 20, labels)
|
|
pyplot.show()
|
|
torch.save((Y, files[:limit]), "tsne_output.pth")
|
|
|
|
|
|
# Uses the results from the calculation above to create a **massive** pdf plot that shows 1/8 size images on the tsne
|
|
# spectrum.
|
|
def plot_results_as_image_graph():
|
|
Y, files = torch.load('tsne_output.pth')
|
|
fig, ax = pyplot.subplots()
|
|
fig.set_size_inches(200,200,forward=True)
|
|
ax.update_datalim(np.column_stack([Y[:, 0], Y[:, 1]]))
|
|
ax.autoscale()
|
|
|
|
for b in tqdm(range(Y.shape[0])):
|
|
im = pyplot.imread(files[b])
|
|
im = OffsetImage(im, zoom=1/8)
|
|
ab = AnnotationBbox(im, (Y[b, 0], Y[b, 1]), xycoords='data', frameon=False)
|
|
ax.add_artist(ab)
|
|
ax.scatter(Y[:, 0], Y[:, 1])
|
|
|
|
pyplot.savefig('tsne.pdf')
|
|
|
|
|
|
if __name__ == "__main__":
|
|
#run_tsne()
|
|
plot_results_as_image_graph() |