<jupyter_start><jupyter_text>Prepare data<jupyter_code>from os.path import join
from pathlib import Path

import torch
import torch.nn.functional as F
import pytorch_lightning as pl

from torch import nn
from torch.utils.data import random_split
from torch.utils.data import DataLoader
from pytorch_lightning.metrics.functional import accuracy
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint

from audio_loader.features.raw_audio import WindowedAudio
from audio_loader.features.log_mel import WindowedLogMel
from audio_loader.ground_truth.timit import TimitGroundTruth
from audio_loader.samplers.dynamic_sampler import DynamicSamplerFromGt
from audio_loader.dl_frontends.pytorch.fill_ram import get_dataset_dynamic_size<jupyter_output><empty_output><jupyter_text>Dataloader lightning<jupyter_code>class TimitMELDataModule(pl.LightningDataModule):
    
    def __init__(self, data_dir, batch_size):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
    
    def prepare_data(self):
        self.timit_gt = TimitGroundTruth(self.data_dir, with_silences=False)
        self.mel_feature_processor = WindowedLogMel(400, 160, 16000, 40, normalize=False)#, delta_orders=[2])
        self.mel_sampler = DynamicSamplerFromGt([self.mel_feature_processor], self.timit_gt)
        self.original_train_dataset, self.collate_func = get_dataset_dynamic_size(self.mel_sampler, "train")
        self.test_dataset, self.collate_func = get_dataset_dynamic_size(self.mel_sampler, "test")
        
    def setup(self, stage=None):
        if stage == 'fit' or stage is None:
            self.val_nb_samples = round(len(self.original_train_dataset)/100)
            self.train_nb_samples = len(self.original_train_dataset) - self.val_nb_samples
            self.train_dataset, self.val_dataset = random_split(
                self.original_train_dataset,
                [self.train_nb_samples, self.val_nb_samples]
            )
        
        if stage == 'test' or stage is None:
            return
            

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True,
                          collate_fn=self.collate_func,
                          drop_last=True)

    def val_dataloader(self):
        #return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False,
        #                  collate_fn=self.collate_func,
        #                  drop_last=False)
        return DataLoader(self.test_dataset, self.batch_size , shuffle=False,
                          collate_fn=self.collate_func,
                          drop_last=False)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, self.batch_size , shuffle=False,
                          collate_fn=self.collate_func,
                          drop_last=False)
# the data
mel_timit = TimitMELDataModule(join(Path.home(), "data/kaggle_TIMIT"), 16)<jupyter_output><empty_output><jupyter_text>Model definition<jupyter_code>class lit_mel_model(pl.LightningModule):
    def __init__(self, feature_size):
        """Init all parameters.
        
        feature_size: int
            size of the expected features for the forward step
        """
        super().__init__()
        self.feature_size = feature_size
        
        self.layer_1_grus = nn.GRU(
            feature_size, 550, 5,
            bidirectional=True,
            batch_first=True,
            dropout=0.2
        )
        
        self.bn_fwd = nn.BatchNorm1d(550)
        self.bn_bwd = nn.BatchNorm1d(550)
        self.layer_2_dense = torch.nn.Linear(1100, 128)
        self.bn_layer_2 = nn.BatchNorm1d(128)
        self.act_layer_2 = nn.LeakyReLU(0.1) # in pytorch kaldi it is softmax
        
        self.layer_3_dense = torch.nn.Linear(128, 58)
        
    def forward(self, x):
        """Forward of the model over the data."""
        batch_size = x.batch_sizes[0]
        # shape: (num_layers*directions, batch_size, hidden_size?)
        h_0 = torch.zeros(5*2, batch_size, 550, device="cuda")
        output, h_n = self.layer_1_grus(x, h_0)

        fwd_h = h_n.view(5, 2, batch_size, 550)[-1, 0]
        bwd_h = h_n.view(5, 2, batch_size, 550)[-1, 1]
    
        fwd_h = self.bn_fwd(fwd_h.view(batch_size, 550))
        bwd_h = self.bn_bwd(bwd_h.view(batch_size, 550))

        h = torch.cat((fwd_h, bwd_h), 1)
        dense1 = self.bn_layer_2(self.act_layer_2(self.layer_2_dense(h)))
        return self.layer_3_dense(dense1)
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.0004)
        return optimizer
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x.to('cuda'))
        _, y = torch.stack(y).max(dim=1)
        
        loss = F.cross_entropy(y_hat, y)
        self.log('train_loss', loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x.to('cuda'))
        _, target = torch.stack(y).max(dim=1)
        _, pred = y_hat.max(dim=1)

        loss = F.cross_entropy(y_hat, target)
        self.log('val_loss', loss)
        self.log('val_accuracy', accuracy(pred+1, target+1))
        return loss
    
    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x.to('cuda'))
        _, target = torch.stack(y).max(dim=1)
        _, pred = y_hat.max(dim=1)
        
        loss = F.cross_entropy(y_hat, target)
        self.log('test_loss', loss)
        self.log('test_accuracy', accuracy(pred+1, target+1))
        return loss
# init model
model = lit_mel_model(40*1) # log MEl band
model.to('cuda')<jupyter_output><empty_output><jupyter_text>Train model<jupyter_code># trainer definition
trainer = pl.Trainer(
    callbacks=[
        EarlyStopping(monitor='val_loss', patience=10, mode="min")
    ],
    checkpoint_callback=ModelCheckpoint(save_top_k=5, monitor="val_loss", mode="min"),
    progress_bar_refresh_rate=1000,
    gpus=1, auto_select_gpus=True,
    precision=16,
    max_epochs=100
)
trainer.fit(model, mel_timit)

mel_timit.setup("test")
trainer.test(model, mel_timit.test_dataloader())<jupyter_output>| Name          | Type        | Params
----------------------------------------------
0 | layer_1_grus  | GRU         | 23 M  
1 | bn_fwd        | BatchNorm1d | 1 K   
2 | bn_bwd        | BatchNorm1d | 1 K   
3 | layer_2_dense | Linear      | 140 K 
4 | bn_layer_2    | BatchNorm1d | 256   
5 | act_layer_2   | LeakyReLU   | 0     
6 | layer_3_dense | Linear      | 7 K
{'fraction_non_alphanumeric': 0.10662867426514697, 'fraction_numerical': 0.022345530893821235, 'mean_word_length': 2.646254784034992, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '284159', 'license': 'MIT', 'n_tokens_mistral': 2356, 'n_tokens_neox': 2241, 'n_words': 400, 'path': 'timit_model_bidir_fbank-without_silence.ipynb', 'repo_name': 'vroger11/audio_loader', 'size': '6668'}