extract_mfcc.py

import pickle
import time
import torch
import torch.nn as nn
import torchaudio
from glob import glob
from nemo.collections.asr.modules import AudioToMFCCPreprocessor
from tqdm import tqdm
import torch.multiprocessing as mp
from itertools import islice
import numpy as np

class EnhancedAudioPreprocessor(nn.Module):
    def __init__(self,
                 sample_rate=16000,
                 n_fft=512,
                 win_length=400,
                 hop_length=160,
                 f_min=20,
                 f_max=7600,
                 n_mels=80,
                 n_mfcc=None,
                 norm_signal=False,
                 do_preemph=True,
                 spec_norm='mn'):
        super().__init__()
        
        self.preprocessor = AudioToMFCCPreprocessor(
            sample_rate=sample_rate,
            n_window_size=win_length,
            n_window_stride=hop_length,
            n_fft=n_fft,
            lowfreq=f_min,
            highfreq=f_max,
            n_mels=n_mels,
            n_mfcc=n_mfcc if n_mfcc is not None else n_mels,
            window_size=None,
            window_stride=None,
        )
        
        self.spec_norm = spec_norm
        if spec_norm == 'mn':
            self.spec_norm_fn = self.mean_norm
        elif spec_norm == 'mvn':
            self.spec_norm_fn = self.mean_var_norm
        elif spec_norm == 'bn':
            self.spec_norm_fn = nn.BatchNorm1d(n_mels)
        else:
            self.spec_norm_fn = nn.Identity()

    def mean_norm(self, x):
        return x - torch.mean(x, dim=-1, keepdim=True)
    
    def mean_var_norm(self, x):
        mean = torch.mean(x, dim=-1, keepdim=True)
        std = torch.std(x, dim=-1, keepdim=True)
        return (x - mean) / (std + 1e-8)

    def forward(self, x, length=None):
        if length is None:
            length = torch.tensor([x.shape[1]] * x.shape[0])
        
        x, length = self.preprocessor(input_signal=x, length=length)
        
        x = x.transpose(1, 2)  # NeMo outputs (B, T, D), we need (B, D, T)
        if isinstance(self.spec_norm_fn, (nn.BatchNorm1d, nn.Identity)):
            x = self.spec_norm_fn(x)
        else:
            x = self.spec_norm_fn(x)
        
        return x, length

def split_list(lst, n):
    """리스트를 n개의 부분으로 나누는 함수"""
    k, m = divmod(len(lst), n)
    return [lst[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n)]

def chunk_list(lst, chunk_size):
    """리스트를 지정된 크기의 청크로 나누는 함수"""
    for i in range(0, len(lst), chunk_size):
        yield lst[i:i + chunk_size]

def create_model_and_preprocessor(device):
    """모델과 전처리기를 생성하는 함수"""
    preprocessor = EnhancedAudioPreprocessor(
        sample_rate=16000,
        n_fft=512,
        win_length=400,
        hop_length=160,
        f_min=20,
        f_max=7600,
        n_mels=80,
        norm_signal=True,
        do_preemph=True,
        spec_norm='mvn'
    ).to(device)
    return preprocessor

def process_batch(file_paths, device_id, result_queue):
    """각 프로세스에서 실행될 함수"""
    device = torch.device('cuda')
    
    # 각 프로세스에서 모델과 전처리기 초기화
    preprocessor = create_model_and_preprocessor(device)
    
    results = []
    for path in file_paths:
        waveform, sample_rate = torchaudio.load(path)
        waveform = waveform.to(device)
        
        if waveform.shape[0] > 1:
            waveform = torch.mean(waveform, dim=0, keepdim=True)
            
        with torch.no_grad():
            processed_signal, processed_length = preprocessor(waveform)
            processed_signal = processed_signal.cpu().numpy()
            processed_length = processed_length.cpu().numpy()
        
        del waveform
        torch.cuda.empty_cache()
        
        results.append((path, processed_signal, processed_length))
    
    result_queue.put(results)

def main():
    num_workers = 12  # 워커(모델) 수 5900X 였기 때문에 그러함
    FILE_PATH = "audio/*"
    
    # 전체 파일 목록 가져오기
    audio_files = glob(FILE_PATH)
    
    # 파일 목록을 워커 수만큼 분할
    split_files = split_list(audio_files, num_workers)
    
    # 멀티프로세싱 설정
    mp.set_start_method('spawn', force=True)
    
    # 결과 큐 생성
    result_queue = mp.Queue()
    
    # 프로세스 생성 및 시작
    processes = []
    for worker_id in range(num_workers):
        p = mp.Process(
            target=process_batch,
            args=(split_files[worker_id], worker_id, result_queue)
        )
        p.start()
        processes.append(p)
    
    all_results = []
    for _ in range(num_workers):
        all_results.extend(result_queue.get())
    
    for p in processes:
        p.join()
    
    return all_results

if __name__ == '__main__':
    start_time = time.time()
    results = main()
    print("elapsed time", time.time() - start_time)
    #results를 파일로 저장
    with open("results.pkl", "wb") as f:
        pickle.dump(results, f)