| | import os |
| | from glob import glob |
| | from pathlib import Path |
| | import torch |
| | import logging |
| | import argparse |
| | import torch |
| | import numpy as np |
| | from sklearn.cluster import KMeans, MiniBatchKMeans |
| | import tqdm |
| | logging.basicConfig(level=logging.INFO) |
| | logger = logging.getLogger(__name__) |
| | import time |
| | import random |
| |
|
| | def train_cluster(in_dir, n_clusters, use_minibatch=True, verbose=False): |
| |
|
| | logger.info(f"Loading features from {in_dir}") |
| | features = [] |
| | nums = 0 |
| | for path in tqdm.tqdm(in_dir.glob("*.soft.pt")): |
| | features.append(torch.load(path).squeeze(0).numpy().T) |
| | |
| | features = np.concatenate(features, axis=0) |
| | print(nums, features.nbytes/ 1024**2, "MB , shape:",features.shape, features.dtype) |
| | features = features.astype(np.float32) |
| | logger.info(f"Clustering features of shape: {features.shape}") |
| | t = time.time() |
| | if use_minibatch: |
| | kmeans = MiniBatchKMeans(n_clusters=n_clusters,verbose=verbose, batch_size=4096, max_iter=80).fit(features) |
| | else: |
| | kmeans = KMeans(n_clusters=n_clusters,verbose=verbose).fit(features) |
| | print(time.time()-t, "s") |
| |
|
| | x = { |
| | "n_features_in_": kmeans.n_features_in_, |
| | "_n_threads": kmeans._n_threads, |
| | "cluster_centers_": kmeans.cluster_centers_, |
| | } |
| | print("end") |
| |
|
| | return x |
| |
|
| |
|
| | if __name__ == "__main__": |
| |
|
| | parser = argparse.ArgumentParser() |
| | parser.add_argument('--dataset', type=Path, default="./dataset/44k", |
| | help='path of training data directory') |
| | parser.add_argument('--output', type=Path, default="logs/44k", |
| | help='path of model output directory') |
| |
|
| | args = parser.parse_args() |
| |
|
| | checkpoint_dir = args.output |
| | dataset = args.dataset |
| | n_clusters = 10000 |
| |
|
| | ckpt = {} |
| | for spk in os.listdir(dataset): |
| | if os.path.isdir(dataset/spk): |
| | print(f"train kmeans for {spk}...") |
| | in_dir = dataset/spk |
| | x = train_cluster(in_dir, n_clusters, verbose=False) |
| | ckpt[spk] = x |
| |
|
| | checkpoint_path = checkpoint_dir / f"kmeans_{n_clusters}.pt" |
| | checkpoint_path.parent.mkdir(exist_ok=True, parents=True) |
| | torch.save( |
| | ckpt, |
| | checkpoint_path, |
| | ) |
| |
|
| |
|
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| |
|
| |
|
| |
|
