feat: train index

app.py

@@ -1,4 +1,8 @@
 import os
+import traceback
+
+import numpy as np
+from sklearn.cluster import MiniBatchKMeans
 
 os.environ["PYTORCH_JIT"] = "0v"
 
@@ -7,6 +11,7 @@ import gradio as gr
 import zipfile
 import tempfile
 import shutil
+import faiss
 from glob import glob
 from infer.modules.train.preprocess import PreProcess
 from infer.modules.train.extract.extract_f0_rmvpe import FeatureInput
@@ -193,6 +198,66 @@ def download_weight(exp_dir: str) -> str:
     return "assets/weights/%s.pth" % name
 
 
+def train_index(exp_dir: str) -> str:
+    feature_dir = "%s/3_feature768" % (exp_dir)
+    if not os.path.exists(feature_dir):
+        raise gr.Error("Please extract features first.")
+    listdir_res = list(os.listdir(feature_dir))
+    if len(listdir_res) == 0:
+        raise gr.Error("Please extract features first.")
+    npys = []
+    for name in sorted(listdir_res):
+        phone = np.load("%s/%s" % (feature_dir, name))
+        npys.append(phone)
+    big_npy = np.concatenate(npys, 0)
+    big_npy_idx = np.arange(big_npy.shape[0])
+    np.random.shuffle(big_npy_idx)
+    big_npy = big_npy[big_npy_idx]
+    if big_npy.shape[0] > 2e5:
+        print("Trying doing kmeans %s shape to 10k centers." % big_npy.shape[0])
+        try:
+            big_npy = (
+                MiniBatchKMeans(
+                    n_clusters=10000,
+                    verbose=True,
+                    batch_size=256 * 8,
+                    compute_labels=False,
+                    init="random",
+                )
+                .fit(big_npy)
+                .cluster_centers_
+            )
+        except:
+            info = traceback.format_exc()
+            print(info)
+            raise gr.Error(info)
+
+    np.save("%s/total_fea.npy" % exp_dir, big_npy)
+    n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
+    print("%s,%s" % (big_npy.shape, n_ivf))
+    index = faiss.index_factory(768, "IVF%s,Flat" % n_ivf)
+    # index = faiss.index_factory(256if version19=="v1"else 768, "IVF%s,PQ128x4fs,RFlat"%n_ivf)
+    print("training")
+    index_ivf = faiss.extract_index_ivf(index)  #
+    index_ivf.nprobe = 1
+    index.train(big_npy)
+    faiss.write_index(
+        index,
+        "%s/trained_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe),
+    )
+    print("adding")
+    batch_size_add = 8192
+    for i in range(0, big_npy.shape[0], batch_size_add):
+        index.add(big_npy[i : i + batch_size_add])
+    faiss.write_index(
+        index,
+        "%s/added_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe),
+    )
+    print("built added_IVF%s_Flat_nprobe_%s.index" % (n_ivf, index_ivf.nprobe))
+
+    return "%s/added_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe)
+
+
 def download_expdir(exp_dir: str) -> str:
     shutil.make_archive(exp_dir, "zip", exp_dir)
     return f"{exp_dir}.zip"
@@ -206,7 +271,7 @@ def restore_expdir(zip: str) -> str:
 
 with gr.Blocks() as app:
     # allow user to manually select the experiment directory
-    exp_dir = gr.Textbox(label="Experiment directory", visible=True, interactive=True)
+    exp_dir = gr.Textbox(label="Experiment directory (don't touch it unless you know what you are doing)", visible=True, interactive=True)
 
     with gr.Tabs():
         with gr.Tab(label="New / Restore"):
@@ -244,8 +309,10 @@ with gr.Blocks() as app:
         with gr.Tab(label="Train"):
             with gr.Row():
                 train_btn = gr.Button(value="Train", variant="primary")
-            with gr.Row():
                 latest_model = gr.File(label="Latest checkpoint")
+            with gr.Row():
+                train_index_btn = gr.Button(value="Train index", variant="primary")
+                trained_index = gr.File(label="Trained index")
 
         with gr.Tab(label="Download"):
             with gr.Row():
@@ -278,6 +345,12 @@ with gr.Blocks() as app:
         outputs=[latest_model],
     )
 
+    train_index_btn.click(
+        fn=train_index,
+        inputs=[exp_dir],
+        outputs=[trained_index],
+    )
+
     download_weight_btn.click(
         fn=download_weight,
         inputs=[exp_dir],

infer/modules/train/train.py

@@ -200,8 +200,7 @@ def run(rank, n_gpus, hps, logger: logging.Logger, state):
         )
         state["global_step"] = (epoch_str - 1) * len(train_loader)
         print("loaded", epoch_str)
-        # epoch_str = 1
-        # global_step = 0
+        epoch_str += 1
     except:  # 如果首次不能加载，加载pretrain
         # traceback.print_exc()
         epoch_str = 1
@@ -248,7 +247,7 @@ def run(rank, n_gpus, hps, logger: logging.Logger, state):
     scaler = GradScaler(enabled=hps.train.fp16_run)
 
     cache = []
-    trained = 0
+    saved = 0
     for epoch in range(epoch_str, hps.train.epochs + 1):
         if rank == 0:
             train_and_evaluate(
@@ -283,8 +282,9 @@ def run(rank, n_gpus, hps, logger: logging.Logger, state):
         scheduler_g.step()
         scheduler_d.step()
 
-        trained += 1
-        if trained >= 10:
+        if epoch % hps.save_every_epoch == 0 and rank == 0:
+            saved += 1
+        if saved >= 2:
             break