Allow BYOL resnet playground to produce a latent dict

codes/scripts/byol_resnet_playground.py

@@ -103,10 +103,30 @@ def get_latent_for_img(model, img):
     return latent
 
 
+def produce_latent_dict(model):
+    batch_size = 32
+    num_workers = 4
+    dataloader = get_image_folder_dataloader(batch_size, num_workers)
+    id = 0
+    paths = []
+    latents = []
+    for batch in tqdm(dataloader):
+        hq = batch['hq'].to('cuda')
+        model(hq)
+        l = layer_hooked_value.cpu().split(1, dim=0)
+        latents.extend(l)
+        paths.extend(batch['HQ_path'])
+        id += batch_size
+        if id > 1000:
+            print("Saving checkpoint..")
+            torch.save((latents, paths), 'results.pth')
+            id = 0
+
+
 def find_similar_latents(model, compare_fn=structural_euc_dist):
     global layer_hooked_value
 
-    img = 'F:\\4k6k\\datasets\\ns_images\\adrianna\\analyze\\analyze_xx\\adrianna_xx.jpg'
+    img = 'F:\\4k6k\\datasets\\ns_images\\imagesets\\1024_test\\80692045.jpg.jpg'
     #img = 'F:\\4k6k\\datasets\\ns_images\\adrianna\\analyze\\analyze_xx\\nicky_xx.jpg'
     output_path = '../../results/byol_resnet_similars'
     os.makedirs(output_path, exist_ok=True)
@@ -153,4 +173,5 @@ if __name__ == '__main__':
     register_hook(model, 'avgpool')
 
     with torch.no_grad():
-        find_similar_latents(model, structural_euc_dist)
+        #find_similar_latents(model, structural_euc_dist)
+        produce_latent_dict(model)

codes/scripts/tsne_torch.py

@@ -0,0 +1,263 @@
+#
+#  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()