更新多说话人训练webui，提供cli命令 / Update the multi-speaker training WebUI and provide CLI commands

train/data_utils.py

@@ -40,7 +40,8 @@ class TextAudioLoaderMultiNSFsid(torch.utils.data.Dataset):
         # spec_length = wav_length // hop_length
         audiopaths_and_text_new = []
         lengths = []
-        for audiopath, text, pitch, pitchf, dv in self.audiopaths_and_text:
+        for record in self.audiopaths_and_text:
+            audiopath, text, pitch, pitchf, dv = record[:5]
             if self.min_text_len <= len(text) and len(text) <= self.max_text_len:
                 audiopaths_and_text_new.append([audiopath, text, pitch, pitchf, dv])
                 lengths.append(os.path.getsize(audiopath) // (3 * self.hop_length))
@@ -248,7 +249,8 @@ class TextAudioLoader(torch.utils.data.Dataset):
         # spec_length = wav_length // hop_length
         audiopaths_and_text_new = []
         lengths = []
-        for audiopath, text, dv in self.audiopaths_and_text:
+        for record in self.audiopaths_and_text:
+            audiopath, text, dv = record[:3]
             if self.min_text_len <= len(text) and len(text) <= self.max_text_len:
                 audiopaths_and_text_new.append([audiopath, text, dv])
                 lengths.append(os.path.getsize(audiopath) // (3 * self.hop_length))

train/preprocess.py

@@ -10,6 +10,7 @@ n_p = int(sys.argv[3])
 exp_dir = sys.argv[4]
 noparallel = sys.argv[5] == "True"
 per = float(sys.argv[6])
+manifest_path = sys.argv[7] if len(sys.argv) > 7 else ""
 import traceback
 
 import librosa
@@ -21,6 +22,7 @@ from infer.audio import load_audio
 from train.dataset.slicer2 import Slicer
 from i18n.i18n import I18nAuto
 from tools.progress import should_report
+from tools.multispeaker import ManifestError, load_manifest
 
 i18n = I18nAuto()
 
@@ -57,19 +59,19 @@ class PreProcess:
         os.makedirs(self.gt_wavs_dir, exist_ok=True)
         os.makedirs(self.wavs16k_dir, exist_ok=True)
 
-    def norm_write(self, tmp_audio, idx0, idx1):
+    def norm_write(self, tmp_audio, output_key, idx1):
         tmp_max = np.abs(tmp_audio).max()
         if not np.isfinite(tmp_max) or tmp_max <= 0 or tmp_max > 2.5:
             println(
                 i18n("[数据切分][跳过] 无效或异常音频片段：%s_%s | 峰值：%s")
-                % (idx0, idx1, tmp_max)
+                % (output_key, idx1, tmp_max)
             )
             return False
         tmp_audio = (tmp_audio / tmp_max * (self.max * self.alpha)) + (
             1 - self.alpha
         ) * tmp_audio
         wavfile.write(
-            "%s/%s_%s.wav" % (self.gt_wavs_dir, idx0, idx1),
+            "%s/%s_%s.wav" % (self.gt_wavs_dir, output_key, idx1),
             self.sr,
             tmp_audio.astype(np.float32),
         )
@@ -77,13 +79,13 @@ class PreProcess:
             tmp_audio, orig_sr=self.sr, target_sr=16000
         ).astype(np.float32)
         wavfile.write(
-            "%s/%s_%s.wav" % (self.wavs16k_dir, idx0, idx1),
+            "%s/%s_%s.wav" % (self.wavs16k_dir, output_key, idx1),
             16000,
             audio_16k,
         )
         return True
 
-    def pipeline(self, path, idx0, total):
+    def pipeline(self, path, output_key, progress_index, total):
         try:
             audio = load_audio(path, self.sr)
             # zero phased digital filter cause pre-ringing noise...
@@ -98,17 +100,17 @@ class PreProcess:
                     i += 1
                     if len(audio[start:]) > self.tail * self.sr:
                         tmp_audio = audio[start : start + int(self.per * self.sr)]
-                        self.norm_write(tmp_audio, idx0, idx1)
+                        self.norm_write(tmp_audio, output_key, idx1)
                         idx1 += 1
                     else:
                         tmp_audio = audio[start:]
                         idx1 += 1
                         break
-                self.norm_write(tmp_audio, idx0, idx1)
-            if should_report(idx0, total):
+            self.norm_write(tmp_audio, output_key, idx1)
+            if should_report(progress_index, total):
                 println(
                     i18n("[数据切分] 进度：%s/%s | %s")
-                    % (idx0 + 1, total, os.path.basename(path))
+                    % (progress_index + 1, total, os.path.basename(path))
                 )
             return True
         except Exception:
@@ -121,8 +123,8 @@ class PreProcess:
     def pipeline_mp(self, infos):
         success = 0
         failed = 0
-        for path, idx0, total in infos:
-            if self.pipeline(path, idx0, total):
+        for path, output_key, progress_index, total in infos:
+            if self.pipeline(path, output_key, progress_index, total):
                 success += 1
             else:
                 failed += 1
@@ -134,10 +136,13 @@ class PreProcess:
 
     def pipeline_mp_inp_dir(self, inp_root, n_p):
         try:
-            names = sorted(os.listdir(inp_root))
+            names = sorted(
+                name for name in os.listdir(inp_root)
+                if os.path.isfile(os.path.join(inp_root, name))
+            )
             total = len(names)
             infos = [
-                ("%s/%s" % (inp_root, name), idx, total)
+                ("%s/%s" % (inp_root, name), str(idx), idx, total)
                 for idx, name in enumerate(names)
             ]
             worker_count = max(n_p, 1)
@@ -162,11 +167,59 @@ class PreProcess:
         except Exception:
             println(i18n("[数据切分][失败] %s") % traceback.format_exc())
 
+    def pipeline_mp_manifest(self, manifest_entries, n_p):
+        infos = [
+            (entry["path"], entry["output_key"], idx, len(manifest_entries))
+            for idx, entry in enumerate(manifest_entries)
+        ]
+        total = len(infos)
+        worker_count = max(n_p, 1)
+        worker_count = min(worker_count, max(total, 1))
+        println(
+            i18n("[数据切分] 多说话人待处理：%s | 进程数：%s")
+            % (total, worker_count)
+        )
+        if noparallel:
+            for i in range(worker_count):
+                self.pipeline_mp(infos[i::worker_count])
+            return
+        ps = []
+        for i in range(worker_count):
+            p = multiprocessing.Process(
+                target=self.pipeline_mp, args=(infos[i::worker_count],)
+            )
+            ps.append(p)
+            p.start()
+        for p in ps:
+            p.join()
+
 
 def preprocess_trainset(inp_root, sr, n_p, exp_dir, per):
     pp = PreProcess(sr, exp_dir, per)
     println(i18n("[数据切分] 开始"))
-    pp.pipeline_mp_inp_dir(inp_root, n_p)
+    if manifest_path:
+        try:
+            manifest = load_manifest(exp_dir)
+            for output_dir in (pp.gt_wavs_dir, pp.wavs16k_dir):
+                for name in os.listdir(output_dir):
+                    if name.startswith("ms") and name.endswith(".wav"):
+                        try:
+                            os.remove(os.path.join(output_dir, name))
+                        except OSError:
+                            pass
+            pp.pipeline_mp_manifest(manifest["entries"], n_p)
+        except ManifestError as error:
+            println(i18n(error.key) % error.values)
+            raise
+    else:
+        for output_dir in (pp.gt_wavs_dir, pp.wavs16k_dir):
+            for name in os.listdir(output_dir):
+                if name.startswith("ms") and name.endswith(".wav"):
+                    try:
+                        os.remove(os.path.join(output_dir, name))
+                    except OSError:
+                        pass
+        pp.pipeline_mp_inp_dir(inp_root, n_p)
     println(i18n("[数据切分] 完成"))
 
 

train/process_ckpt.py

@@ -1,15 +1,33 @@
 import os
 import sys
 import traceback
+import json
 from collections import OrderedDict
 
 import torch
 
 from i18n.i18n import I18nAuto
+from tools.file_io import read_text
 
 i18n = I18nAuto()
 
 
+def normalize_speaker_info(speaker_info):
+    result = []
+    seen = set()
+    for item in speaker_info or []:
+        try:
+            speaker_id = int(item["id"])
+            speaker_name = str(item["name"])
+        except (KeyError, TypeError, ValueError):
+            continue
+        if speaker_id < 0 or speaker_id > 109 or not speaker_name or speaker_id in seen:
+            continue
+        seen.add(speaker_id)
+        result.append({"id": speaker_id, "name": speaker_name})
+    return sorted(result, key=lambda item: item["id"])
+
+
 def savee(ckpt, sr, if_f0, name, epoch, version, hps):
     try:
         opt = OrderedDict()
@@ -42,6 +60,9 @@ def savee(ckpt, sr, if_f0, name, epoch, version, hps):
         opt["sr"] = sr
         opt["f0"] = if_f0
         opt["version"] = version
+        speaker_info = normalize_speaker_info(getattr(hps, "speaker_info", []))
+        if speaker_info:
+            opt["speaker_info"] = speaker_info
         torch.save(opt, "assets/weights/%s.pth" % name)
         return i18n("成功")
     except:
@@ -63,6 +84,11 @@ def show_info(path):
 
 def extract_small_model(path, name, sr, if_f0, info, version):
     try:
+        speaker_info = []
+        config_path = os.path.join(os.path.dirname(os.path.abspath(path)), "config.json")
+        if os.path.isfile(config_path):
+            config_data = json.loads(read_text(config_path))
+            speaker_info = normalize_speaker_info(config_data.get("speaker_info", []))
         ckpt = torch.load(path, map_location="cpu")
         if "model" in ckpt:
             ckpt = ckpt["model"]
@@ -185,6 +211,8 @@ def extract_small_model(path, name, sr, if_f0, info, version):
         opt["version"] = version
         opt["sr"] = sr
         opt["f0"] = int(if_f0)
+        if speaker_info:
+            opt["speaker_info"] = speaker_info
         torch.save(opt, "assets/weights/%s.pth" % name)
         return i18n("成功")
     except:
@@ -218,6 +246,8 @@ def merge(path1, path2, alpha1, sr, f0, info, name, version):
 
         ckpt1 = torch.load(path1, map_location="cpu")
         ckpt2 = torch.load(path2, map_location="cpu")
+        speaker_info1 = normalize_speaker_info(ckpt1.get("speaker_info", []))
+        speaker_info2 = normalize_speaker_info(ckpt2.get("speaker_info", []))
         cfg = ckpt1["config"]
         if "model" in ckpt1:
             ckpt1 = extract(ckpt1)
@@ -255,6 +285,8 @@ def merge(path1, path2, alpha1, sr, f0, info, name, version):
         opt["f0"] = 1 if f0 == i18n("是") else 0
         opt["version"] = version
         opt["info"] = info
+        if speaker_info1 and speaker_info1 == speaker_info2:
+            opt["speaker_info"] = speaker_info1
         torch.save(opt, "assets/weights/%s.pth" % name)
         return i18n("成功")
     except:

train/train.py

@@ -106,6 +106,27 @@ class EpochRecorder:
         return f"[{current_time}] | ({elapsed_time_str})"
 
 
+def load_pretrained_generator(model, path):
+    target = model.module if hasattr(model, "module") else model
+    saved_state = torch.load(path, map_location="cpu")["model"]
+    current_state = target.state_dict()
+    embedding_key = "emb_g.weight"
+    if embedding_key in saved_state and embedding_key in current_state:
+        saved_embedding = saved_state[embedding_key]
+        current_embedding = current_state[embedding_key]
+        if saved_embedding.shape != current_embedding.shape:
+            compatible = (
+                saved_embedding.dim() == current_embedding.dim()
+                and saved_embedding.shape[1:] == current_embedding.shape[1:]
+            )
+            if compatible:
+                expanded = current_embedding.clone()
+                rows = min(saved_embedding.shape[0], current_embedding.shape[0])
+                expanded[:rows].copy_(saved_embedding[:rows])
+                saved_state[embedding_key] = expanded
+    return target.load_state_dict(saved_state)
+
+
 def main():
     n_gpus = torch.cuda.device_count()
     single_cuda = torch.cuda.is_available() and n_gpus == 1
@@ -242,18 +263,7 @@ def run(rank, n_gpus, hps, logger, use_ddp):
         if hps.pretrainG != "":
             if rank == 0:
                 logger.info(i18n("已加载生成器预训练模型：%s") % hps.pretrainG)
-            if hasattr(net_g, "module"):
-                logger.info(
-                    net_g.module.load_state_dict(
-                        torch.load(hps.pretrainG, map_location="cpu")["model"]
-                    )
-                )  ##测试不加载优化器
-            else:
-                logger.info(
-                    net_g.load_state_dict(
-                        torch.load(hps.pretrainG, map_location="cpu")["model"]
-                    )
-                )  ##测试不加载优化器
+            logger.info(load_pretrained_generator(net_g, hps.pretrainG))
         if hps.pretrainD != "":
             if rank == 0:
                 logger.info(i18n("已加载判别器预训练模型：%s") % hps.pretrainD)

train/train_index.py

@@ -3,6 +3,7 @@ import platform
 import sys
 import traceback
 import glob
+import json
 
 import faiss
 import numpy as np
@@ -18,6 +19,7 @@ exp_name = sys.argv[1]
 version = sys.argv[2]
 outside_index_root = sys.argv[3]
 n_cpu = int(sys.argv[4])
+index_mode = sys.argv[5] if len(sys.argv) > 5 else "auto"
 exp_dir = os.path.join("logs", exp_name)
 feature_dir = os.path.join(
     exp_dir, "3_feature256" if version == "v1" else "3_feature768"
@@ -41,14 +43,20 @@ def newest_index(pattern):
     return max(paths, key=os.path.getmtime) if paths else ""
 
 
-def link_added_index(added_path):
+def speaker_scope(message, speaker_id):
+    if speaker_id is None:
+        return message
+    return "%s：%s | %s" % (i18n("说话人ID（0~109）"), speaker_id, message)
+
+
+def link_added_index(added_path, speaker_id=None):
     added_name = os.path.basename(added_path)
     try:
         os.makedirs(outside_index_root, exist_ok=True)
         source = os.path.abspath(added_path)
         outside_root = os.path.abspath(outside_index_root)
         if os.path.commonpath([source, outside_root]) == outside_root:
-            log(i18n("[索引训练] 外部索引链接已存在：%s") % source)
+            log(speaker_scope(i18n("[索引训练] 外部索引链接已存在：%s") % source, speaker_id))
             return
         target = os.path.abspath(
             os.path.join(outside_index_root, "%s_%s" % (exp_name, added_name))
@@ -56,7 +64,7 @@ def link_added_index(added_path):
         if os.path.lexists(target):
             try:
                 if os.path.samefile(source, target):
-                    log(i18n("[索引训练] 外部索引链接已存在：%s") % target)
+                    log(speaker_scope(i18n("[索引训练] 外部索引链接已存在：%s") % target, speaker_id))
                     return
             except (FileNotFoundError, OSError):
                 pass
@@ -65,110 +73,177 @@ def link_added_index(added_path):
             os.link(source, target)
         else:
             os.symlink(source, target)
-        log(i18n("[索引训练] 已链接索引到外部目录：%s") % outside_index_root)
+        log(speaker_scope(i18n("[索引训练] 已链接索引到外部目录：%s") % outside_index_root, speaker_id))
     except Exception:
         log(
-            i18n("[索引训练][失败] 无法链接索引到外部目录：%s\n%s")
-            % (outside_index_root, traceback.format_exc())
+            speaker_scope(
+                i18n("[索引训练][失败] 无法链接索引到外部目录：%s\n%s")
+                % (outside_index_root, traceback.format_exc()),
+                speaker_id,
+            )
         )
 
 
-existing_trained_path = newest_index(
-    os.path.join(
-        exp_dir,
-        "trained_IVF*_Flat_nprobe_*_%s_%s.index" % (exp_name, version),
-    )
-)
-existing_added_path = newest_index(
-    os.path.join(exp_dir, "added_IVF*_Flat_nprobe_*_%s_%s.index" % (exp_name, version))
-)
-if not existing_added_path:
-    existing_added_path = newest_index(
-        os.path.join(
-            outside_index_root,
-            "%s_*added_IVF*_Flat_nprobe_*_%s_%s.index"
-            % (exp_name, exp_name, version),
-        )
-    )
-if existing_added_path:
-    if existing_trained_path:
-        log(
-            i18n("[索引训练][跳过] trained索引已存在：%s")
-            % os.path.basename(existing_trained_path)
-        )
-    log(
-        i18n("[索引训练][跳过] added索引已存在：%s")
-        % os.path.basename(existing_added_path)
-    )
-    link_added_index(existing_added_path)
-    raise SystemExit(0)
-
 if not os.path.isdir(feature_dir) or not os.listdir(feature_dir):
     log(i18n("[索引训练][失败] 请先进行特征提取"))
     raise SystemExit(1)
 
-features = []
-for name in sorted(os.listdir(feature_dir)):
-    features.append(np.load(os.path.join(feature_dir, name)))
-
-big_npy = np.concatenate(features, 0)
-big_npy = big_npy[np.random.permutation(big_npy.shape[0])]
-if big_npy.shape[0] > 200000:
-    log(i18n("[索引训练] 正在将%s条特征聚类为10000个中心") % big_npy.shape[0])
+manifest_path = os.path.join(exp_dir, "multispeaker_manifest.json")
+manifest_by_key = {}
+if index_mode != "single" and os.path.isfile(manifest_path):
     try:
-        big_npy = MiniBatchKMeans(
-            n_clusters=10000,
-            verbose=False,
-            batch_size=256 * n_cpu,
-            compute_labels=False,
-            init="random",
-        ).fit(big_npy).cluster_centers_
-    except Exception:
-        log(i18n("[索引训练][失败] 聚类失败，将使用原始特征继续\n%s") % traceback.format_exc())
+        with open(manifest_path, "r", encoding="utf8") as file:
+            manifest = json.load(file)
+        entries = manifest.get("entries", []) if isinstance(manifest, dict) else []
+        manifest_by_key = {
+            str(entry["output_key"]): int(entry["speaker_id"])
+            for entry in entries
+        }
+    except (KeyError, TypeError, ValueError, OSError, json.JSONDecodeError):
+        manifest_by_key = {}
 
-n_ivf = max(1, min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39))
-log(i18n("[索引训练] 特征形状：%s | IVF数量：%s") % (big_npy.shape, n_ivf))
-if existing_trained_path:
-    trained_path = existing_trained_path
-    index = faiss.read_index(trained_path)
-    index_ivf = faiss.extract_index_ivf(index)
-    index_ivf.nprobe = 1
-    log(
-        i18n("[索引训练][跳过] trained索引已存在：%s")
-        % os.path.basename(trained_path)
-    )
+feature_paths = [
+    os.path.join(feature_dir, name)
+    for name in sorted(os.listdir(feature_dir))
+    if name.lower().endswith(".npy")
+]
+feature_groups = {}
+if manifest_by_key:
+    for path in feature_paths:
+        stem = os.path.splitext(os.path.basename(path))[0]
+        output_key = stem if stem in manifest_by_key else stem.rsplit("_", 1)[0]
+        speaker_id = manifest_by_key.get(output_key)
+        if speaker_id is not None:
+            feature_groups.setdefault(speaker_id, []).append(path)
 else:
-    index = faiss.index_factory(
-        256 if version == "v1" else 768, "IVF%s,Flat" % n_ivf
-    )
-    index_ivf = faiss.extract_index_ivf(index)
-    index_ivf.nprobe = 1
-    trained_path = os.path.join(
+    feature_groups[None] = feature_paths
+
+if not feature_groups or not any(feature_groups.values()):
+    log(i18n("[索引训练][失败] 请先进行特征提取"))
+    raise SystemExit(1)
+
+
+def train_one_speaker(speaker_id, paths):
+    scope = lambda message: speaker_scope(message, speaker_id)
+    features = [np.load(path) for path in paths]
+    big_npy = np.concatenate(features, 0)
+    big_npy = big_npy[np.random.permutation(big_npy.shape[0])]
+    suffix = "" if speaker_id is None else "_spkid%s" % speaker_id
+    total_path = os.path.join(exp_dir, "total_fea%s.npy" % suffix)
+    np.save(total_path, big_npy)
+
+    trained_pattern = os.path.join(
         exp_dir,
-        "trained_IVF%s_Flat_nprobe_%s_%s_%s.index"
-        % (n_ivf, index_ivf.nprobe, exp_name, version),
+        "trained_IVF*_Flat_nprobe_*_%s_%s%s.index"
+        % (exp_name, version, suffix),
     )
-    log(i18n("[索引训练] 正在训练索引"))
-    index.train(big_npy)
-    faiss.write_index(index, trained_path)
-
-log(i18n("[索引训练] 正在写入特征向量"))
-starts = list(range(0, big_npy.shape[0], 8192))
-for batch_index, start in enumerate(starts):
-    index.add(big_npy[start : start + 8192])
-    if should_report(batch_index, len(starts), 10):
-        log(
-            i18n("[索引训练] 写入进度：%s/%s")
-            % (batch_index + 1, len(starts))
+    added_pattern = os.path.join(
+        exp_dir,
+        "added_IVF*_Flat_nprobe_*_%s_%s%s.index"
+        % (exp_name, version, suffix),
+    )
+    existing_trained_path = newest_index(trained_pattern)
+    existing_added_path = newest_index(added_pattern)
+    if not existing_added_path:
+        existing_added_path = newest_index(
+            os.path.join(
+                outside_index_root,
+                "%s_*added_IVF*_Flat_nprobe_*_%s_%s%s.index"
+                % (exp_name, exp_name, version, suffix),
+            )
         )
+    if existing_added_path:
+        if existing_trained_path:
+            log(
+                scope(
+                    i18n("[索引训练][跳过] trained索引已存在：%s")
+                    % os.path.basename(existing_trained_path)
+                )
+            )
+        log(
+            scope(
+                i18n("[索引训练][跳过] added索引已存在：%s")
+                % os.path.basename(existing_added_path)
+            )
+        )
+        link_added_index(existing_added_path, speaker_id)
+        return
 
-added_name = "added_IVF%s_Flat_nprobe_%s_%s_%s.index" % (
-    n_ivf,
-    index_ivf.nprobe,
-    exp_name,
-    version,
-)
-added_path = os.path.join(exp_dir, added_name)
-faiss.write_index(index, added_path)
-log(i18n("[索引训练] 成功构建索引：%s") % added_name)
-link_added_index(added_path)
+    if big_npy.shape[0] > 200000:
+        log(
+            scope(
+                i18n("[索引训练] 正在将%s条特征聚类为10000个中心")
+                % big_npy.shape[0]
+            )
+        )
+        try:
+            big_npy = MiniBatchKMeans(
+                n_clusters=10000,
+                verbose=False,
+                batch_size=256 * n_cpu,
+                compute_labels=False,
+                init="random",
+            ).fit(big_npy).cluster_centers_
+        except Exception:
+            log(
+                scope(
+                    i18n("[索引训练][失败] 聚类失败，将使用原始特征继续\n%s")
+                    % traceback.format_exc()
+                )
+            )
+
+    n_ivf = max(1, min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39))
+    log(scope(i18n("[索引训练] 特征形状：%s | IVF数量：%s") % (big_npy.shape, n_ivf)))
+    if existing_trained_path:
+        trained_path = existing_trained_path
+        index = faiss.read_index(trained_path)
+        index_ivf = faiss.extract_index_ivf(index)
+        index_ivf.nprobe = 1
+        log(
+            scope(
+                i18n("[索引训练][跳过] trained索引已存在：%s")
+                % os.path.basename(trained_path)
+            )
+        )
+    else:
+        index = faiss.index_factory(
+            256 if version == "v1" else 768, "IVF%s,Flat" % n_ivf
+        )
+        index_ivf = faiss.extract_index_ivf(index)
+        index_ivf.nprobe = 1
+        trained_path = os.path.join(
+            exp_dir,
+            "trained_IVF%s_Flat_nprobe_%s_%s_%s%s.index"
+            % (n_ivf, index_ivf.nprobe, exp_name, version, suffix),
+        )
+        log(scope(i18n("[索引训练] 正在训练索引")))
+        index.train(big_npy)
+        faiss.write_index(index, trained_path)
+
+    log(scope(i18n("[索引训练] 正在写入特征向量")))
+    starts = list(range(0, big_npy.shape[0], 8192))
+    for batch_index, start in enumerate(starts):
+        index.add(big_npy[start : start + 8192])
+        if should_report(batch_index, len(starts), 10):
+            log(
+                scope(
+                    i18n("[索引训练] 写入进度：%s/%s")
+                    % (batch_index + 1, len(starts))
+                )
+            )
+
+    added_name = "added_IVF%s_Flat_nprobe_%s_%s_%s%s.index" % (
+        n_ivf,
+        index_ivf.nprobe,
+        exp_name,
+        version,
+        suffix,
+    )
+    added_path = os.path.join(exp_dir, added_name)
+    faiss.write_index(index, added_path)
+    log(scope(i18n("[索引训练] 成功构建索引：%s") % added_name))
+    link_added_index(added_path, speaker_id)
+
+
+for speaker_id in sorted(feature_groups, key=lambda value: -1 if value is None else value):
+    train_one_speaker(speaker_id, feature_groups[speaker_id])

train/utils.py

@@ -105,14 +105,25 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
     saved_state_dict = checkpoint_dict["model"]
     if hasattr(model, "module"):
         state_dict = model.module.state_dict()
-    else:
-        state_dict = model.state_dict()
-    new_state_dict = {}
+    else:
+        state_dict = model.state_dict()
+    new_state_dict = {}
+    embedding_resized = False
     for k, v in state_dict.items():  # 模型需要的shape
-        try:
-            new_state_dict[k] = saved_state_dict[k]
-            if saved_state_dict[k].shape != state_dict[k].shape:
-                logger.warning(
+        try:
+            new_state_dict[k] = saved_state_dict[k]
+            if saved_state_dict[k].shape != state_dict[k].shape:
+                if (
+                    k == "emb_g.weight"
+                    and saved_state_dict[k].dim() == state_dict[k].dim()
+                    and saved_state_dict[k].shape[1:] == state_dict[k].shape[1:]
+                ):
+                    new_state_dict[k] = v.clone()
+                    rows = min(saved_state_dict[k].shape[0], state_dict[k].shape[0])
+                    new_state_dict[k][:rows].copy_(saved_state_dict[k][:rows])
+                    embedding_resized = True
+                    continue
+                logger.warning(
                     "shape-%s-mismatch|need-%s|get-%s",
                     k,
                     state_dict[k].shape,
@@ -131,9 +142,9 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
 
     iteration = checkpoint_dict["iteration"]
     learning_rate = checkpoint_dict["learning_rate"]
-    if (
-        optimizer is not None and load_opt == 1
-    ):  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
+    if (
+        optimizer is not None and load_opt == 1 and not embedding_resized
+    ):  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
         #   try:
         optimizer.load_state_dict(checkpoint_dict["optimizer"])
     #   except: