forked from mrq/DL-Art-School
92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
|
import torch
|
||
|
|
import numpy as np
|
||
|
|
from random import shuffle
|
||
|
|
from tqdm import tqdm
|
||
|
|
import random
|
||
|
|
import os
|
||
|
|
import time
|
||
|
|
import torch
|
||
|
|
from matplotlib import pyplot as plt
|
||
|
|
from pykeops.torch import LazyTensor
|
||
|
|
|
||
|
|
use_cuda = False
|
||
|
|
dtype = torch.float32
|
||
|
|
device_id = 'cpu'
|
||
|
|
|
||
|
|
|
||
|
|
def load_vectors():
|
||
|
|
""" Will need to be modified per-data type you are loading. """
|
||
|
|
all_files = torch.load('/y/separated/large_mel_cheaters_linux.pth')
|
||
|
|
os.makedirs('/y/separated/randomly_sampled_cheaters', exist_ok=True)
|
||
|
|
vecs = []
|
||
|
|
print("Gathering vectors..")
|
||
|
|
j = 0
|
||
|
|
for f in tqdm(all_files):
|
||
|
|
vs=torch.tensor(np.load(f)['arr_0'])
|
||
|
|
for k in range(4):
|
||
|
|
vecs.append(vs[0,:,random.randint(0,vs.shape[-1]-1)])
|
||
|
|
if len(vecs) >= 1000000:
|
||
|
|
vecs = torch.stack(vecs, dim=0)
|
||
|
|
torch.save(vecs, f'/y/separated/randomly_sampled_cheaters/{j}.pth')
|
||
|
|
j += 1
|
||
|
|
vecs = []
|
||
|
|
vecs = [torch.stack(vecs, dim=0)]
|
||
|
|
for i in range(j):
|
||
|
|
vecs.append(torch.load(f'/y/separated/randomly_sampled_cheaters/{i}.pth'))
|
||
|
|
vecs = torch.cat(vecs, dim=0)
|
||
|
|
torch.save(vecs, '/y/separated/randomly_sampled_cheaters/combined.pth')
|
||
|
|
|
||
|
|
def k_means(x, K, Niter=10, verbose=True):
|
||
|
|
"""Implements Lloyd's algorithm for the Euclidean metric.
|
||
|
|
Thanks to https://www.kernel-operations.io/keops/_auto_tutorials/kmeans/plot_kmeans_torch.html
|
||
|
|
"""
|
||
|
|
|
||
|
|
start = time.time()
|
||
|
|
N, D = x.shape # Number of samples, dimension of the ambient space
|
||
|
|
|
||
|
|
c = x[:K, :].clone() # Simplistic initialization for the centroids
|
||
|
|
|
||
|
|
x_i = LazyTensor(x.view(N, 1, D)) # (N, 1, D) samples
|
||
|
|
c_j = LazyTensor(c.view(1, K, D)) # (1, K, D) centroids
|
||
|
|
|
||
|
|
# K-means loop:
|
||
|
|
# - x is the (N, D) point cloud,
|
||
|
|
# - cl is the (N,) vector of class labels
|
||
|
|
# - c is the (K, D) cloud of cluster centroids
|
||
|
|
for i in tqdm(range(Niter)):
|
||
|
|
|
||
|
|
# E step: assign points to the closest cluster -------------------------
|
||
|
|
D_ij = ((x_i - c_j) ** 2).sum(-1) # (N, K) symbolic squared distances
|
||
|
|
cl = D_ij.argmin(dim=1).long().view(-1) # Points -> Nearest cluster
|
||
|
|
|
||
|
|
# M step: update the centroids to the normalized cluster average: ------
|
||
|
|
# Compute the sum of points per cluster:
|
||
|
|
c.zero_()
|
||
|
|
c.scatter_add_(0, cl[:, None].repeat(1, D), x)
|
||
|
|
|
||
|
|
# Divide by the number of points per cluster:
|
||
|
|
Ncl = torch.bincount(cl, minlength=K).type_as(c).view(K, 1)
|
||
|
|
c /= Ncl # in-place division to compute the average
|
||
|
|
|
||
|
|
if verbose: # Fancy display -----------------------------------------------
|
||
|
|
if use_cuda:
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
end = time.time()
|
||
|
|
print(
|
||
|
|
f"K-means for the Euclidean metric with {N:,} points in dimension {D:,}, K = {K:,}:"
|
||
|
|
)
|
||
|
|
print(
|
||
|
|
"Timing for {} iterations: {:.5f}s = {} x {:.5f}s\n".format(
|
||
|
|
Niter, end - start, Niter, (end - start) / Niter
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
return cl, c
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
#load_vectors()
|
||
|
|
vecs = torch.load('/y/separated/randomly_sampled_cheaters/combined.pth')
|
||
|
|
cl, c = k_means(vecs, 8192, 50)
|
||
|
|
torch.save((cl, c), '/y/separated/randomly_sampled_cheaters/k_means_clusters.pth')
|