{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Learnable Time-Frequency Representations\nThis assignment was about training models to learn different time-frequency representations of audio. I chose to work with the Rainforest dataset because I'm familiar with the dataset from a previous Kaggle competition and because I thought it was fun to use a bioacoustics dataset for the assignment.","metadata":{}},{"cell_type":"code","source":"import math\nfrom pathlib import Path\nimport random\nimport uuid\nimport time\n\nimport IPython.display as ipd\nimport librosa\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torchaudio\nfrom tqdm.notebook import tqdm\n\ntorchaudio.set_audio_backend('sox_io')","metadata":{"execution":{"iopub.status.busy":"2021-06-08T20:27:58.170210Z","iopub.execute_input":"2021-06-08T20:27:58.170683Z","iopub.status.idle":"2021-06-08T20:28:01.331749Z","shell.execute_reply.started":"2021-06-08T20:27:58.170609Z","shell.execute_reply":"2021-06-08T20:28:01.330814Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's define some paths to the Rainforest dataset","metadata":{}},{"cell_type":"code","source":"input_data = Path('/kaggle/input')\noutput_data = Path('/kaggle/working')\n\nrainforest_data = input_data / 'rfcx-species-audio-detection'\n\ntrain_data = rainforest_data / 'train'\ntest_data = rainforest_data / 'test'\n\ndf = pd.read_csv(rainforest_data / 'train_tp.csv')","metadata":{"execution":{"iopub.status.busy":"2021-06-08T20:28:16.737775Z","iopub.execute_input":"2021-06-08T20:28:16.738123Z","iopub.status.idle":"2021-06-08T20:28:16.755705Z","shell.execute_reply.started":"2021-06-08T20:28:16.738087Z","shell.execute_reply":"2021-06-08T20:28:16.754952Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The dataframe describing the Rainforest dataset has the following format:","metadata":{}},{"cell_type":"code","source":"df.head()","metadata":{"execution":{"iopub.status.busy":"2021-06-08T20:28:21.472845Z","iopub.execute_input":"2021-06-08T20:28:21.473176Z","iopub.status.idle":"2021-06-08T20:28:21.495596Z","shell.execute_reply.started":"2021-06-08T20:28:21.473145Z","shell.execute_reply":"2021-06-08T20:28:21.494574Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"For each row in the dataframe we are going to extract a short segment of audio of the file corresponding to `recording_id`. We are going to center the extracted segment in the middle of `t_min` and `t_max`, and all extracted segments will be of the same lenght, for simplicity. We are going to do the segmentation in advance and save the extracted and normalized segments as tensors, for faster loading when we train the models.\nWe specify the wanted length and sample rate of the extracted segments, as well as the encoder we want to use:","metadata":{}},{"cell_type":"code","source":"SAMPLE_RATE = 22050\nCLIP_LEN_SECONDS = 1.\nCLIP_LEN_SAMPLES = int(SAMPLE_RATE * CLIP_LEN_SECONDS)\nN_FFT = 300\nENCODER = torchaudio.transforms.Spectrogram(n_fft=N_FFT)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T20:28:23.928362Z","iopub.execute_input":"2021-06-08T20:28:23.928691Z","iopub.status.idle":"2021-06-08T20:28:23.992077Z","shell.execute_reply.started":"2021-06-08T20:28:23.928661Z","shell.execute_reply":"2021-06-08T20:28:23.991315Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"And we create a directory where we will save the extracted segments","metadata":{}},{"cell_type":"code","source":"waveform_tensors = output_data / 'waveform-tensors'\n\nPath.mkdir(waveform_tensors)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T20:28:26.140687Z","iopub.execute_input":"2021-06-08T20:28:26.141013Z","iopub.status.idle":"2021-06-08T20:28:26.145764Z","shell.execute_reply.started":"2021-06-08T20:28:26.140977Z","shell.execute_reply":"2021-06-08T20:28:26.144446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"weights_dir = output_data / 'weights'\n\nPath.mkdir(weights_dir, exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T20:28:27.079616Z","iopub.execute_input":"2021-06-08T20:28:27.080016Z","iopub.status.idle":"2021-06-08T20:28:27.084726Z","shell.execute_reply.started":"2021-06-08T20:28:27.079967Z","shell.execute_reply":"2021-06-08T20:28:27.083532Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"`get_normalized_segment` takes a filepath and a timestamp as input and returns a normalized waveform of length `CLIP_LEN_SECONDS` centered around the provided timestamp. We loop through the entire dataframe and find all the interesting segments, we store them as torch tensors in the `waveform_tensors` directory.","metadata":{}},{"cell_type":"code","source":"def get_normalized_segment(fpath, mid_segment_timestamp):\n    audio, sr = torchaudio.load(fpath)\n    audio = audio.squeeze()\n    resampler = torchaudio.transforms.Resample(sr, SAMPLE_RATE)\n    audio = resampler(audio)\n    start_idx = int(sr * (mid_segment_timestamp - CLIP_LEN_SECONDS / 2))\n    end_idx = start_idx + CLIP_LEN_SAMPLES\n    if start_idx < 0:\n        start_idx = 0\n        end_idx = start_idx + CLIP_LEN_SAMPLES\n    elif end_idx > len(audio) - 1:\n        end_idx = len(audio) - 1\n        start_idx = end_idx - CLIP_LEN_SAMPLES\n    seg = audio[start_idx:end_idx]\n    seg -= torch.mean(seg)\n    seg /= torch.max(torch.abs(seg))\n    return seg\n\nfor row_idx, row in tqdm(df.iterrows(), total=df.shape[0]):\n    fpath = train_data / (row['recording_id'] + '.flac')\n    mid_segment_timestamp = (row['t_min'] + row['t_max']) / 2.\n    segment = get_normalized_segment(fpath, mid_segment_timestamp)\n    assert segment.shape == torch.Size([CLIP_LEN_SAMPLES])\n    fname = str(uuid.uuid4()) + '.pt'\n    torch.save(segment, waveform_tensors / fname)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T20:28:29.238940Z","iopub.execute_input":"2021-06-08T20:28:29.239273Z","iopub.status.idle":"2021-06-08T20:39:52.673376Z","shell.execute_reply.started":"2021-06-08T20:28:29.239241Z","shell.execute_reply":"2021-06-08T20:39:52.672513Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We create a class `WaveformDataset` to conveniently load data as we need it. We also create a few helper functions for vizualizing the data.","metadata":{}},{"cell_type":"code","source":"class WaveformDataset(Dataset):\n    def __init__(self, fpaths, encoder):\n        self.fpaths = fpaths\n        self.encoder = encoder\n\n    def __len__(self):\n        return len(self.fpaths)\n\n    def __getitem__(self, idx):\n        return torch.load(self.fpaths[idx])\n        \n    def show_sample(self, idx=None, waveform=None):\n        if waveform is None:\n            assert idx is not None\n            waveform = torch.load(self.fpaths[idx])\n        self.plot_waveform(waveform)\n        self.plot_encoded(waveform)\n        self.display_audio(waveform)\n        \n    def plot_waveform(self, waveform):\n        plt.figure()\n        plt.title('Waveform')\n        plt.plot(waveform.detach().numpy())\n    \n    def plot_encoded(self, waveform):\n        encoded_waveform = self.encoder(waveform)\n        plt.figure()\n        plt.title('Encoded waveform')\n        plt.imshow(encoded_waveform.detach().numpy())\n        \n    def display_audio(self, waveform):\n        ipd.display(ipd.Audio(waveform.detach().numpy(), rate=SAMPLE_RATE))\n\n    def show_random_sample(self):\n        self.show_sample(idx=random.randrange(self.__len__()))","metadata":{"execution":{"iopub.status.busy":"2021-06-08T20:39:55.712700Z","iopub.execute_input":"2021-06-08T20:39:55.713024Z","iopub.status.idle":"2021-06-08T20:39:55.724430Z","shell.execute_reply.started":"2021-06-08T20:39:55.712993Z","shell.execute_reply":"2021-06-08T20:39:55.723560Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's create a dataset and look at a random sample","metadata":{}},{"cell_type":"code","source":"ds = WaveformDataset(list(waveform_tensors.iterdir()), ENCODER)\nds.show_random_sample()","metadata":{"execution":{"iopub.status.busy":"2021-06-08T20:39:58.587511Z","iopub.execute_input":"2021-06-08T20:39:58.587845Z","iopub.status.idle":"2021-06-08T20:39:58.992770Z","shell.execute_reply.started":"2021-06-08T20:39:58.587817Z","shell.execute_reply":"2021-06-08T20:39:58.992061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We now define the model we are going to use. The model takes normalized audio waveforms as input, the first part of the model encodes the input with some waveform -> time-frequency transformation. The second part of the model is the decoder, it consists of transposed 1d convolution blocks that upsamples the time dimension of the activations through the layers and simultanously downsamples the frequency dimension, producing a pure time domain signal","metadata":{}},{"cell_type":"code","source":"class Decoder(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        self.tconv1 = nn.ConvTranspose1d(\n            in_channels=151,\n            out_channels=128,\n            kernel_size=16,\n            stride=4,\n            dilation=4,\n        )\n        \n        self.tconv2 = nn.ConvTranspose1d(\n            in_channels=128,\n            out_channels=64,\n            kernel_size=32,\n            stride=4,\n            dilation=4,\n        )\n        \n        self.tconv3 = nn.ConvTranspose1d(\n            in_channels=64,\n            out_channels=32,\n            kernel_size=32,\n            stride=4,\n            dilation=4,\n        )\n        \n        self.tconv4 = nn.ConvTranspose1d(\n            in_channels=32,\n            out_channels=1,\n            kernel_size=74,\n            stride=2,\n            dilation=1,\n        )\n        \n    def forward(self, x):\n        x = self.tconv1(x)\n        x = F.leaky_relu(x)\n        x = self.tconv2(x)\n        x = F.leaky_relu(x)\n        x = self.tconv3(x)\n        x = F.leaky_relu(x)\n        x = self.tconv4(x)\n        x = x.squeeze()\n        return x\n\nclass EncoderDecoder(nn.Module):\n    def __init__(self, encoder):\n        super().__init__()\n\n        self.encoder = encoder\n        self.decoder = Decoder()\n\n    def forward(self, x):\n        x = self.encoder(x)\n        x = self.decoder(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2021-06-08T21:02:23.354891Z","iopub.execute_input":"2021-06-08T21:02:23.355226Z","iopub.status.idle":"2021-06-08T21:02:23.364948Z","shell.execute_reply.started":"2021-06-08T21:02:23.355188Z","shell.execute_reply":"2021-06-08T21:02:23.364014Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"After trying multiple things, what worked \"best\" was to try to minimize the L1 loss between the encoded input and output. I also tried to minimize the L1 loss of the waveforms directly, but that gave even worse results.","metadata":{}},{"cell_type":"code","source":"def spec_loss(pred, target, encoder, criterion):\n    encoded_pred = encoder(pred)\n    encoded_target = encoder(target)\n    log_pred = torch.log(encoded_pred + 1e-8)\n    log_target = torch.log(encoded_target+ 1e-8)\n    l1 = criterion(log_pred, log_target)\n    return l1","metadata":{"execution":{"iopub.status.busy":"2021-06-08T21:02:25.039475Z","iopub.execute_input":"2021-06-08T21:02:25.039796Z","iopub.status.idle":"2021-06-08T21:02:25.045256Z","shell.execute_reply.started":"2021-06-08T21:02:25.039768Z","shell.execute_reply":"2021-06-08T21:02:25.044420Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The training loop is pretty straightforward, I decrease the learning rate by a factor of 10 when the validation loss stops to decrease.","metadata":{}},{"cell_type":"code","source":"def train(model, encoder, train_ds, val_ds, weights_path, device, batch_size=32, epochs=20):\n\n    train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n    val_dl = DataLoader(val_ds, batch_size=batch_size)\n\n    optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-2)\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=2, verbose=True, factor=0.1)\n\n    criterion = nn.L1Loss()\n\n    model = model.to(device)\n    criterion = criterion.to(device)\n    encoder = encoder.to(device)\n\n    best_val_loss = math.inf\n\n    for epoch in tqdm(range(epochs), desc='Training'):\n        start_time = time.time()\n        model.train()\n        train_loss = 0.\n\n        for waveform_train_batch in train_dl:\n            waveform_train_batch = waveform_train_batch.to(device)\n            waveform_train_preds = model(waveform_train_batch)\n            loss = spec_loss(waveform_train_preds, waveform_train_batch, encoder, criterion)\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n            train_loss += loss.item()\n\n        train_loss /= len(train_dl)\n\n        with torch.no_grad():\n            val_loss = 0.\n            model.eval()\n            for waveform_val_batch in val_dl:\n                waveform_val_batch = waveform_val_batch.to(device)\n                waveform_val_preds = model(waveform_val_batch)\n                loss = spec_loss(waveform_val_preds, waveform_val_batch, encoder, criterion)\n                val_loss += loss.item()\n\n            val_loss /= len(val_dl)\n\n        elapsed = time.time() - start_time\n        print(f'Epoch {epoch+1} (time: {elapsed:.0f}s): train_loss: {train_loss} val_loss: {val_loss}')\n\n        if val_loss < best_val_loss:\n            print(f'Saving new best model at epoch {epoch} (val_loss improved from {best_val_loss} to {val_loss})')\n            torch.save(model, weights_path)\n            best_val_loss = val_loss\n\n        scheduler.step(val_loss)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T21:02:29.192708Z","iopub.execute_input":"2021-06-08T21:02:29.193031Z","iopub.status.idle":"2021-06-08T21:02:29.203618Z","shell.execute_reply.started":"2021-06-08T21:02:29.193001Z","shell.execute_reply":"2021-06-08T21:02:29.202679Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"weights_path = weights_dir / 'weights1.pt'\n\nfpaths = list(waveform_tensors.iterdir())\ntrain_fpaths, val_fpaths = train_test_split(fpaths, test_size=0.33)\n\ntrain_ds = WaveformDataset(train_fpaths, ENCODER)\nval_ds = WaveformDataset(val_fpaths, ENCODER)\n\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\nmodel = EncoderDecoder(ENCODER)\n\ntrain(model, ENCODER, train_ds, val_ds, weights_path, device, epochs=50)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T21:02:30.697415Z","iopub.execute_input":"2021-06-08T21:02:30.697745Z","iopub.status.idle":"2021-06-08T21:23:03.739972Z","shell.execute_reply.started":"2021-06-08T21:02:30.697716Z","shell.execute_reply":"2021-06-08T21:23:03.739110Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can now compare a few audio snippets to their reconstructed counterparts","metadata":{}},{"cell_type":"code","source":"model = torch.load(weights_dir / 'weights1.pt')\nmodel = model.cpu()\nmodel.eval()\n\nval_ds = WaveformDataset(val_fpaths, ENCODER.cpu())","metadata":{"execution":{"iopub.status.busy":"2021-06-08T21:23:41.086811Z","iopub.execute_input":"2021-06-08T21:23:41.087144Z","iopub.status.idle":"2021-06-08T21:23:41.106144Z","shell.execute_reply.started":"2021-06-08T21:23:41.087114Z","shell.execute_reply":"2021-06-08T21:23:41.105181Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x = val_ds[random.randrange(len(val_ds))]\nval_ds.show_sample(waveform=x)\n\nx_reconstruct = model(x.unsqueeze(0))\nval_ds.show_sample(waveform=x_reconstruct)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T21:25:22.151055Z","iopub.execute_input":"2021-06-08T21:25:22.151500Z","iopub.status.idle":"2021-06-08T21:25:22.716619Z","shell.execute_reply.started":"2021-06-08T21:25:22.151460Z","shell.execute_reply":"2021-06-08T21:25:22.715677Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x = val_ds[random.randrange(len(val_ds))]\nval_ds.show_sample(waveform=x)\n\nx_reconstruct = model(x.unsqueeze(0))\nval_ds.show_sample(waveform=x_reconstruct)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T21:25:44.470894Z","iopub.execute_input":"2021-06-08T21:25:44.471206Z","iopub.status.idle":"2021-06-08T21:25:45.022480Z","shell.execute_reply.started":"2021-06-08T21:25:44.471177Z","shell.execute_reply":"2021-06-08T21:25:45.021491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can see (and hear!) that the reconstructed audio does not sound very good, although we can see that the shape of the waveforms are pretty similar. By looking at the spectrograms, we can see that the model manages to localize the birdcall in time, but the dominating frequencies are off.","metadata":{}}]}