HW4: Speaker Prediction

  • goal: learn to use transformer

1. Data

Dataset

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import os
import json
import torch
import random
from pathlib import Path
from torch.utils.data import Dataset
from torch.nn.utils.rnn import pad_sequence

class myDataset(Dataset):
    '''
    数据集
    '''
    def __init__(self, data_dir, segment_len=128):
        self.data_dir = data_dir
        self.segment_len = segment_len
        
        # 从 name 映射到 id
        mapping_path = Path(data_dir) / "mapping.json"
        mapping = json.load(mapping_path.open())
        self.speaker2id = mapping["speaker2id"]
        
        # 加载 metadata
        metadata_path = Path(data_dir) / "metadata.json"
        metadata = json.load(open(metadata_path))["speakers"]
        
        # 数据统计
        self.speaker_num = len(metadata.keys())
        self.data = []
        for speaker in metadata.keys():
            for uttrances in metadata[speaker]:
                self.data.append([uttrances['feature_path'], self.speaker2id[speaker]])
        
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        feat_path, speaker = self.data[index]
        
        # 加载预处理后的 mel-spectrogram
        mel = torch.load(os.path.join(self.data_dir, feat_path))
        
        # 切片处理
        if len(mel) > self.segment_len:
            # 随机获取 start 位置
            start = random.randint(0, len(mel)-self.segment_len)
            mel = torch.FloatTensor(mel[start : start+self.segment_len])
        else:
            mel = torch.FloatTensor(mel)
            
        speaker = torch.FloatTensor([speaker]).long()
        return mel, speaker
    
    def get_speaker_number(self):
        return self.speaker_num

DataLoader

切分：

  • 90% train

  • 10% validation

创建 DataLoader 用于 iterate

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
import torch
from torch.utils.data import DataLoader, random_split
from torch.nn.utils.rnn import pad_sequence

def collate_batch(batch):
    '''
    逐批量的，因此我们需要填充使得长度一致
    '''
    mel, speaker = zip(*batch)
    mel = pad_sequence(mel, batch_first=True, padding_value=-20)
    return mel, torch.FloatTensor(speaker).long()

def get_dataloader(data_dir, batch_size, n_workers):
    '''
    生成 DataLoader
    '''
    dataset = myDataset(data_dir)
    speaker_num = dataset.get_speaker_number()
    trainlen = int(0.9 * len(dataset))
    lengths = [trainlen, len(dataset) - trainlen]
    trainset, validset = random_split(dataset, lengths)
    
    train_loader = DataLoader(
        trainset,
        batch_size=batch_size,
        shuffle=True,
        drop_last=True,
#         num_workers=n_workers,
        pin_memory=True,
        collate_fn=collate_batch,
    )
    
    valid_loader = DataLoader(
        validset,
        batch_size=batch_size,
#         num_workers=n_workers,
        drop_last=True,
        pin_memory=True,
        collate_fn=collate_batch,
    )
    
    return train_loader, valid_loader, speaker_num

2. Model

  • TransformerEncoderLayer:

    • Base transformer encoder layer in Attention Is All You Need
    • Parameters:

      • d_model: the number of expected features of the input (required).

      • nhead: the number of heads of the multiheadattention models (required).

      • dim_feedforward: the dimension of the feedforward network model (default=2048).

      • dropout: the dropout value (default=0.1).

      • activation: the activation function of intermediate layer, relu or gelu (default=relu).

  • TransformerEncoder:

    • TransformerEncoder is a stack of N transformer encoder layers
    • Parameters:

      • encoder_layer: an instance of the TransformerEncoderLayer() class (required).

      • num_layers: the number of sub-encoder-layers in the encoder (required).

      • norm: the layer normalization component (optional).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import torch
import torch.nn as nn
import torch.nn.functional as F

class Classifier(nn.Module):
    def __init__(self, d_model=80, n_spks=600, dropout=0.1):
        super().__init__()
        
        # 从 输入维度 映射到 d_model 维度
        self.prenet = nn.Linear(40, d_model)
        # TODO:
        #   Change Transformer to Conformer.
        #   https://arxiv.org/abs/2005.08100
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, dim_feedforward=256, nhead=2
        )
        # self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=2)
        
        # 从 d_model 维度映射到 speaker_num 维度
        self.pred_layer = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.ReLU(),
            nn.Linear(d_model, n_spks)
        )
    
    def forward(self, mels):
        '''
        args:
          mels: (batch size, length, 40)
        return:
          out: (batch size, n_spks)
        '''
        # out: (batch size, length, d_model)
        out = self.prenete(mels)
        # out: (length, batch size, d_model)
        out = out.permute(1, 0, 2)
        # 编码层
        out = self.encoder_layer(out)
        # out: (batch size, length, d_model)
        out = out.transpose(0, 1)
        # 平均池化
        stats = out.mean(dim=1)
        
        # out: (batch, n_spks)
        out = self.pred_layer(stats)
        return out

3. Learning Rate Schedule

  • 对于 transformer 架构，学习率的设计与以往不同。

  • warmup 很重要。

    • 起初，lr 设置 0

    • lr 从 0 逐渐线性增长到 initial lr

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import math

import torch
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR


def get_cosine_schedule_with_warmup(
    optimizer: Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    num_cycles: float = 0.5,
    last_epoch: int = -1,
):
    '''
    创建一个学习率的时间表，学习率随着余弦函数之间的值而减小。
     - 初始lr在优化器中设置为0，在预热期间，它在0和之间线性增加
     - 在优化器中设置初始lr。
     
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
          The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
          The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
          The total number of training steps.
        num_cycles (:obj:`float`, `optional`, defaults to 0.5):
          The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
          following a half-cosine).
        last_epoch (:obj:`int`, `optional`, defaults to -1):
          The index of the last epoch when resuming training.

    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    '''
    
    def lr_lambda(current_step):
        # Warmup
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        
        # Decadence
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        
        return max(
            0.0,
            0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))
        )
    
    return LambdaLR(optimizer, lr_lambda, last_epoch)


4. Model Function

model forward function

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import torch

def model_fn(batch, model, criterion, device):
    """
    forward a bacth through the model
    """
    mels, labels = batch
    mels, labels = mels.to(device), labels.to(device)
    
    outs = model(mels)
    
    loss = criterion(outs, labels)
    preds = outs.argmax(1)
    accuracy = torch.mean((preds == labels).float())
    
    return loss, accuracy

5. Validate

计算在 valid set 上的 accuracy

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from tqdm import tqdm
import torch

def valid(dataloader, model, criterion, device):
    """
    计算在 valid set 上的 accuracy
    """
    model.eval()
    running_loss = 0.0
    running_accuracy = 0.0
    pbar = tqdm(total=len(dataloader.dataset), ncols=0, desc='Valid', unit=' uttr')
    
    for i, batch in enumerate(dataloader):
        with torch.no_grad():
            loss, accuracy = model_fn(batch, model, criterion, device)
            running_loss += loss.item()
            running_accuracy += accuracy.item()
            
        pbar.updatea(dataloader.batch_size)
        pbar.set_postfix(
            loss=f'{running_loss / (i+1)  :.2f}',
            accuracy=f'{running_accuracy / (i+1)  :.2f}',
        )
    
    pbar.close()
    model.train()
    
    return running_accuracy / len(dataloader)


6. Main Function

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
from tqdm import tqdm

import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import DataLoader, random_split

def parase_args():
    '''参数'''
    config = {
        'data_dir': './Dataset',
        'save_path': 'model.ckpt',
        'batch_size': 32,
        'n_workers': 8,
        'valid_steps': 2000,
        'warmup_steps': 1000,
        'total_steps': 70000,
        'save_steps': 10000,
    }
    
    return config


def main(
    data_dir,
    save_path,
    batch_size,
    n_workers,
    valid_steps,
    warmup_steps,
    total_steps,
    save_steps
):
    '''
    Main function
    '''
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"[Info]: Use {device} now!")
    
    train_loader, valid_loader, speaker_num = get_dataloader(data_dir, batch_size, n_workers)
    train_iterator = iter(train_loader)
    print(f"[Info]: Finish loading data!", flush=True)
    
    model = Classifier(n_spks=speaker_num).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = Adam(model.parameters(), lr=1e-3)
    scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)
    print(f"[Info]: Finish creating model!", flush=True)
    
    best_accuracy = -1.0
    best_state_dict = None
    
    pbar = tqdm(total=valid_steps, ncols=0, desc='Train', unit=' step')
    
    for step in range(total_steps):
        # Get data
        try:
            batch = next(train_iterator)
        except StopIteration:
            train_iterator = iter(train_loader)
            batch = next(train_iterator)
            
        loss, accuracy = model_fn(batch, model, criterion, device)
        batch_loss = loss.item()
        batch_accuracy = accuracy.item()
        
        # Update model
        loss.backward()
        optimizer.step()
        scheduler.step()
        optimizer.zero_gard()
        
        # Log
        pbar.update()
        pbar.set_postfix(
            loss=f"{batch_loss:.2f}",
            accuracy=f"{batch_accuracy:.2f}",
            step=step+1,
        )
        
        # Do validation
        if (step+1) % valid_steps == 0:
            pbar.close()
            # keep best state
            valid_accuracy = valid(valid_loader, model, criterion, device)
            if valid_accuracy > best_accuracy:
                best_accuracy = valid_accuracy
                best_state_dict = model.state_dict()
                
            pbar = tqdm(total=valid_steps, ncols=0, desc='Train', unit=' step')
        
        # Save the best model so far
        if (step+1) % save_steps == 0 and best_state_dict is not None:
            torch.save(best_state_dict, save_path)
            pbar.write(f"Step {step+1}, best model saved. (accuracy={best_accuracy:.4f})")
            
    pbar.close()
    
if __name__ == '__main__':
    main(**parase_args())

[Info]: Use cpu now!
[Info]: Finish loading data!
[Info]: Finish creating model!


Train:   0% 0/2000 [00:00<?, ? step/s]

7. Inference

Dataset of inference

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import os
import json
import torch
from pathlib import Path
from torch.utils.data import Dataset


class InferenceDataset(Dataset):
    def __init__(self, data_dir):
        testdata_path = Path(data_dir) / "testdata.json"
        metadata = json.load(testdata_path.open())
        self.data_dir = data_dir
        self.data = metadata['utterances']
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        utterance = self.data[index]
        feat_path = utterance['feature_path']
        mel = torch.load(os.path.json(self.data_dir, feat_path))
        
        return feat_path, mel
    
    
def inference_collate_batch(batch):
    '''Collate a batch of data.'''
    feat_path, mels = zip(*batch)
    
    return feat_path, torch.stack(mels)




### Main fuction of Inference


```python
import json
import csv
from pathlib import Path
from tqdm.notebook import tqdm

import torch
from torch.utils.data import DataLoader

def parse_args():
    """arguments"""
    
    config = {
        'data_dir': './Dataset',
        'model_path': './model.ckpt',
        'output_path': './output.csv',
    }
    
    return config

def main(
    data_dir,
    model_path,
    output_path,
):
    """Main Function"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"[Info]: Use {device} now!")

    mapping_path = Path(data_dir) / "mapping.json"
    mapping = json.load(mapping_path.open())

    dataset = InferenceDataset(data_dir)
    dataloader = DataLoader(
        dataset,
        batch_size=1,
        shuffle=False,
        drop_last=False,
        num_workers=8,
        collate_fn=inference_collate_batch,
    )
    print(f"[Info]: Finish loading data!",flush = True)

    speaker_num = len(mapping["id2speaker"])
    model = Classifier(n_spks=speaker_num).to(device)
    model.load_state_dict(torch.load(model_path))
    model.eval()
    print(f"[Info]: Finish creating model!",flush = True)

    results = [["Id", "Category"]]
    for feat_paths, mels in tqdm(dataloader):
        with torch.no_grad():
            mels = mels.to(device)
            outs = model(mels)
            preds = outs.argmax(1).cpu().numpy()
            for feat_path, pred in zip(feat_paths, preds):
                results.append([feat_path, mapping["id2speaker"][str(pred)]])

    with open(output_path, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerows(results)

        
if __name__ == "__main__":
    main(**parse_args())