Upload cluster_all_layers_nonorm_afterrope_group.py

cluster_all_layers_nonorm_afterrope_group.py

@@ -0,0 +1,91 @@
+import os
+import glob
+import torch
+import faiss
+import numpy as np
+import random
+from tqdm import tqdm
+# Data directory
+# data_dir = '../feats_offset'
+
+
+# Parameter settings
+feature_dim = 5120
+num_clusters = 1000  # Number of clusters
+#batch_size  = 10000000
+batch_size  = 1000000
+#batch_size=200000
+niter = 20
+num_tensor_file = int(batch_size/64/576)
+save_folder = f"/sensei-fs/users/wezhao/projects/data/cluster/centroids_faiss_K_c1k_bs1m_iter_{niter}_nonorm_all_layers_afterrope_group"
+os.system("mkdir "+save_folder)
+for layer_idx in range(40):
+    os.system("mkdir "+save_folder+"/"+str(layer_idx))
+    #if layer_idx<=30:
+    #    continue
+    data_dir = '/sensei-fs/users/wezhao/projects/proj-phu/DenseToken/data/key_states_save_13b_all_layers_after_rope/'+str(layer_idx)
+    # Get a list of all .pt files in the directory
+    pt_files = glob.glob(os.path.join(data_dir, '*.pth'))
+
+    print(f"Found {len(pt_files)} .pth files.")
+
+    print("num_tensor_file:",num_tensor_file)
+    tensor_files = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith('.pth') ]
+    random.shuffle(tensor_files)
+
+    # Initialize GPU resources
+    print("Initializing GPU resources...")
+    res = faiss.StandardGpuResources()
+
+    print("Setting up clustering parameters...")
+
+    
+
+    # Data iterator function
+    def data_iterator(tensor_files):
+        # 
+        for i in range(0, len(tensor_files), num_tensor_file):
+            # Read three tensor files at a time
+            tensors = []
+            # import pdb; pdb.set_trace()
+            for j in range(num_tensor_file):
+                if i + j < len(tensor_files):
+                    print("loading " + str(i)+" "+str(j)+tensor_files[i + j])
+                    tensor = torch.load(tensor_files[i + j])
+                    # import pdb; pdb.set_trace()
+                    tensor = tensor.reshape(-1, feature_dim).cpu().numpy().astype(np.float32)
+                    # import pdb; pdb.set_trace()
+                    
+                    tensors.append(tensor)
+
+            if tensors:
+                yield np.concatenate(tensors, axis=0)
+
+    # Fit the MiniBatchKMeans model incrementally
+    count = 0
+    # import pdb; pdb.set_trace()
+    for data_batch in tqdm(data_iterator(tensor_files), desc="Processing batches"):   
+        data_batch = data_batch.reshape(-1, 40, 128)
+        for i in range(40):
+            data = data_batch[:,i,:]
+            # faiss.normalize_L2(data)
+            kmeans = faiss.Kmeans(d=128, k=num_clusters, niter=niter, gpu=True, verbose=True)
+            # print("====")
+            # Train k-means clustering model on GPU
+            print("Training k-means clustering model on GPU...")
+            data = np.ascontiguousarray(data, dtype=np.float32)
+            # import pdb; pdb.set_trace()
+            # faiss.copy_array_to_vector(np.zeros((100000,5120),dtype=np.float32).ravel(), kmeans.centroids)
+            # faiss.vector_to_array(kmeans.centroids)
+            #kmeans.train(data_batch, index)
+            kmeans.train(data)
+            print("k-means training completed.")
+            # Extract centroids
+            print("Extracting centroids...")
+            centroids = kmeans.centroids
+            
+            np.save(save_folder+"/"+f"{layer_idx}/{i}.npy", centroids)
+            #np.save(f"temp/{count}.npy", centroids)
+            print("Centroids saved layer "+str(layer_idx))
+            count += 1
+        break