## Train a Deep Q Learning (DQN) agent on the MSPacman-v0

## Mounting google drive for saving the video

```
from google.colab import drive
drive.mount('/content/gdrive')
```

    Drive already mounted at /content/gdrive; to attempt to forcibly remount, call drive.mount("/content/gdrive", force_remount=True).

## Get the MSPackman ready from Atari

```
import urllib.request
urllib.request.urlretrieve('http://www.atarimania.com/roms/Roms.rar','Roms.rar')
!pip install unrar
!unrar x Roms.rar
!mkdir rars
!mv HC\ ROMS.zip   rars
!mv ROMS.zip  rars
!python -m atari_py.import_roms rars
```

    Collecting unrar
      Downloading unrar-0.4-py3-none-any.whl (25 kB)
    Installing collected packages: unrar
    Successfully installed unrar-0.4

    UNRAR 5.50 freeware      Copyright (c) 1993-2017 Alexander Roshal

    Extracting from Roms.rar

    Extracting  HC ROMS.zip                                                   36%  OK
    Extracting  ROMS.zip                                                      74% 99%  OK
    All OK
    copying adventure.bin from ROMS/Adventure (1980) (Atari, Warren Robinett) (CX2613, CX2613P) (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/adventure.bin
    copying air_raid.bin from ROMS/Air Raid (Men-A-Vision) (PAL) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/air_raid.bin
    copying alien.bin from ROMS/Alien (1982) (20th Century Fox Video Games, Douglas 'Dallas North' Neubauer) (11006) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/alien.bin
    copying amidar.bin from ROMS/Amidar (1982) (Parker Brothers, Ed Temple) (PB5310) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/amidar.bin
    copying assault.bin from ROMS/Assault (AKA Sky Alien) (1983) (Bomb - Onbase) (CA281).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/assault.bin
    copying asterix.bin from ROMS/Asterix (AKA Taz) (07-27-1983) (Atari, Jerome Domurat, Steve Woita) (CX2696) (Prototype).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/asterix.bin
    copying asteroids.bin from ROMS/Asteroids (1981) (Atari, Brad Stewart - Sears) (CX2649 - 49-75163) [no copyright] ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/asteroids.bin
    copying atlantis.bin from ROMS/Atlantis (Lost City of Atlantis) (1982) (Imagic, Dennis Koble) (720103-1A, 720103-1B, IA3203, IX-010-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/atlantis.bin
    copying bank_heist.bin from ROMS/Bank Heist (Bonnie & Clyde, Cops 'n' Robbers, Hold-Up, Roaring 20's) (1983) (20th Century Fox Video Games, Bill Aspromonte) (11012) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/bank_heist.bin
    copying battle_zone.bin from ROMS/Battlezone (1983) (Atari - GCC, Mike Feinstein, Brad Rice) (CX2681) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/battle_zone.bin
    copying beam_rider.bin from ROMS/Beamrider (1984) (Activision - Cheshire Engineering, David Rolfe, Larry Zwick) (AZ-037-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/beam_rider.bin
    copying berzerk.bin from ROMS/Berzerk (1982) (Atari, Dan Hitchens - Sears) (CX2650 - 49-75168) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/berzerk.bin
    copying bowling.bin from ROMS/Bowling (1979) (Atari, Larry Kaplan - Sears) (CX2628 - 6-99842, 49-75117) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/bowling.bin
    copying boxing.bin from ROMS/Boxing - La Boxe (1980) (Activision, Bob Whitehead) (AG-002, CAG-002, AG-002-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/boxing.bin
    copying breakout.bin from ROMS/Breakout - Breakaway IV (Paddle) (1978) (Atari, Brad Stewart - Sears) (CX2622 - 6-99813, 49-75107) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/breakout.bin
    copying carnival.bin from ROMS/Carnival (1982) (Coleco - Woodside Design Associates, Steve 'Jessica Stevens' Kitchen) (2468) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/carnival.bin
    copying centipede.bin from ROMS/Centipede (1983) (Atari - GCC) (CX2676) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/centipede.bin
    copying chopper_command.bin from ROMS/Chopper Command (1982) (Activision, Bob Whitehead) (AX-015, AX-015-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/chopper_command.bin
    copying crazy_climber.bin from ROMS/Crazy Climber (1983) (Atari - Roklan, Joe Gaucher, Alex Leavens) (CX2683) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/crazy_climber.bin
    copying defender.bin from ROMS/Defender (1982) (Atari, Robert C. Polaro, Alan J. Murphy - Sears) (CX2609 - 49-75186) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/defender.bin
    copying demon_attack.bin from ROMS/Demon Attack (Death from Above) (1982) (Imagic, Rob Fulop) (720000-200, 720101-1B, 720101-1C, IA3200, IA3200C, IX-006-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/demon_attack.bin
    copying donkey_kong.bin from ROMS/Donkey Kong (1982) (Coleco - Woodside Design Associates - Imaginative Systems Software, Garry Kitchen) (2451) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/donkey_kong.bin
    copying double_dunk.bin from ROMS/Double Dunk (Super Basketball) (1989) (Atari, Matthew L. Hubbard) (CX26159) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/double_dunk.bin
    copying elevator_action.bin from ROMS/Elevator Action (1983) (Atari, Dan Hitchens) (CX26126) (Prototype) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/elevator_action.bin
    copying enduro.bin from ROMS/Enduro (1983) (Activision, Larry Miller) (AX-026, AX-026-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/enduro.bin
    copying fishing_derby.bin from ROMS/Fishing Derby (1980) (Activision, David Crane) (AG-004) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/fishing_derby.bin
    copying freeway.bin from ROMS/Freeway (1981) (Activision, David Crane) (AG-009, AG-009-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/freeway.bin
    copying frogger.bin from ROMS/Frogger (1982) (Parker Brothers, Ed English, David Lamkins) (PB5300) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/frogger.bin
    copying frostbite.bin from ROMS/Frostbite (1983) (Activision, Steve Cartwright) (AX-031) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/frostbite.bin
    copying galaxian.bin from ROMS/Galaxian (1983) (Atari - GCC, Mark Ackerman, Tom Calderwood, Glenn Parker) (CX2684) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/galaxian.bin
    copying gopher.bin from ROMS/Gopher (Gopher Attack) (1982) (U.S. Games Corporation - JWDA, Sylvia Day, Todd Marshall, Robin McDaniel, Henry Will IV) (VC2001) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/gopher.bin
    copying gravitar.bin from ROMS/Gravitar (1983) (Atari, Dan Hitchens, Mimi Nyden) (CX2685) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/gravitar.bin
    copying hero.bin from ROMS/H.E.R.O. (1984) (Activision, John Van Ryzin) (AZ-036-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/hero.bin
    copying ice_hockey.bin from ROMS/Ice Hockey - Le Hockey Sur Glace (1981) (Activision, Alan Miller) (AX-012, CAX-012, AX-012-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/ice_hockey.bin
    copying jamesbond.bin from ROMS/James Bond 007 (James Bond Agent 007) (1984) (Parker Brothers - On-Time Software, Joe Gaucher, Louis Marbel) (PB5110) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/jamesbond.bin
    copying journey_escape.bin from ROMS/Journey Escape (1983) (Data Age, J. Ray Dettling) (112-006) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/journey_escape.bin
    copying kaboom.bin from ROMS/Kaboom! (Paddle) (1981) (Activision, Larry Kaplan, David Crane) (AG-010, AG-010-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/kaboom.bin
    copying kangaroo.bin from ROMS/Kangaroo (1983) (Atari - GCC, Kevin Osborn) (CX2689) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/kangaroo.bin
    copying keystone_kapers.bin from ROMS/Keystone Kapers - Raueber und Gendarm (1983) (Activision, Garry Kitchen - Ariola) (EAX-025, EAX-025-04I - 711 025-725) (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/keystone_kapers.bin
    copying king_kong.bin from ROMS/King Kong (1982) (Tigervision - Software Electronics Corporation, Karl T. Olinger - Teldec) (7-001 - 3.60001 VE) (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/king_kong.bin
    copying koolaid.bin from ROMS/Kool-Aid Man (Kool Aid Pitcher Man) (1983) (M Network, Stephen Tatsumi, Jane Terjung - Kool Aid) (MT4648) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/koolaid.bin
    copying krull.bin from ROMS/Krull (1983) (Atari, Jerome Domurat, Dave Staugas) (CX2682) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/krull.bin
    copying kung_fu_master.bin from ROMS/Kung-Fu Master (1987) (Activision - Imagineering, Dan Kitchen, Garry Kitchen) (AG-039-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/kung_fu_master.bin
    copying laser_gates.bin from ROMS/Laser Gates (AKA Innerspace) (1983) (Imagic, Dan Oliver) (720118-2A, 13208, EIX-007-04I) (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/laser_gates.bin
    copying lost_luggage.bin from ROMS/Lost Luggage (Airport Mayhem) (1982) (Apollo - Games by Apollo, Larry Minor, Ernie Runyon, Ed Salvo) (AP-2004) [no opening scene] ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/lost_luggage.bin
    copying montezuma_revenge.bin from ROMS/Montezuma's Revenge - Featuring Panama Joe (1984) (Parker Brothers - JWDA, Henry Will IV) (PB5760) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/montezuma_revenge.bin
    copying mr_do.bin from ROMS/Mr. Do! (1983) (CBS Electronics, Ed English) (4L4478) (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/mr_do.bin
    copying ms_pacman.bin from ROMS/Ms. Pac-Man (1983) (Atari - GCC, Mark Ackerman, Glenn Parker) (CX2675) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/ms_pacman.bin
    copying name_this_game.bin from ROMS/Name This Game (Guardians of Treasure) (1983) (U.S. Games Corporation - JWDA, Roger Booth, Sylvia Day, Ron Dubren, Todd Marshall, Robin McDaniel, Wes Trager, Henry Will IV) (VC1007) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/name_this_game.bin
    copying pacman.bin from ROMS/Pac-Man (1982) (Atari, Tod Frye) (CX2646) (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/pacman.bin
    copying phoenix.bin from ROMS/Phoenix (1983) (Atari - GCC, Mike Feinstein, John Mracek) (CX2673) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/phoenix.bin
    copying video_pinball.bin from ROMS/Pinball (AKA Video Pinball) (Zellers).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/video_pinball.bin
    copying pitfall.bin from ROMS/Pitfall! - Pitfall Harry's Jungle Adventure (Jungle Runner) (1982) (Activision, David Crane) (AX-018, AX-018-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/pitfall.bin
    copying pooyan.bin from ROMS/Pooyan (1983) (Konami) (RC 100-X 02) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/pooyan.bin
    copying private_eye.bin from ROMS/Private Eye (1984) (Activision, Bob Whitehead) (AG-034-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/private_eye.bin
    copying qbert.bin from ROMS/Q-bert (1983) (Parker Brothers - Western Technologies, Dave Hampton, Tom Sloper) (PB5360) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/qbert.bin
    copying riverraid.bin from ROMS/River Raid (1982) (Activision, Carol Shaw) (AX-020, AX-020-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/riverraid.bin
    copying road_runner.bin from patched version of ROMS/Road Runner (1989) (Atari - Bobco, Robert C. Polaro) (CX2663) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/road_runner.bin
    copying robotank.bin from ROMS/Robot Tank (Robotank) (1983) (Activision, Alan Miller) (AZ-028, AG-028-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/robotank.bin
    copying seaquest.bin from ROMS/Seaquest (1983) (Activision, Steve Cartwright) (AX-022) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/seaquest.bin
    copying sir_lancelot.bin from ROMS/Sir Lancelot (1983) (Xonox - K-Tel Software - Product Guild, Anthony R. Henderson) (99006, 6220) (PAL).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/sir_lancelot.bin
    copying skiing.bin from ROMS/Skiing - Le Ski (1980) (Activision, Bob Whitehead) (AG-005, CAG-005, AG-005-04) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/skiing.bin
    copying solaris.bin from ROMS/Solaris (The Last Starfighter, Star Raiders II, Universe) (1986) (Atari, Douglas Neubauer, Mimi Nyden) (CX26136) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/solaris.bin
    copying space_invaders.bin from ROMS/Space Invaders (1980) (Atari, Richard Maurer - Sears) (CX2632 - 49-75153) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/space_invaders.bin
    copying star_gunner.bin from ROMS/Stargunner (1983) (Telesys, Alex Leavens) (1005) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/star_gunner.bin
    copying surround.bin from ROMS/Surround (32 in 1) (Bit Corporation) (R320).bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/surround.bin
    copying tennis.bin from ROMS/Tennis - Le Tennis (1981) (Activision, Alan Miller) (AG-007, CAG-007) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/tennis.bin
    copying time_pilot.bin from ROMS/Time Pilot (1983) (Coleco - Woodside Design Associates, Harley H. Puthuff Jr.) (2663) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/time_pilot.bin
    copying trondead.bin from ROMS/TRON - Deadly Discs (TRON Joystick) (1983) (M Network - INTV - APh Technological Consulting, Jeff Ronne, Brett Stutz) (MT5662) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/trondead.bin
    copying tutankham.bin from ROMS/Tutankham (1983) (Parker Brothers, Dave Engman, Dawn Stockbridge) (PB5340) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/tutankham.bin
    copying up_n_down.bin from ROMS/Up 'n Down (1984) (SEGA - Beck-Tech, Steve Beck, Phat Ho) (009-01) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/up_n_down.bin
    copying venture.bin from ROMS/Venture (1982) (Coleco, Joseph Biel) (2457) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/venture.bin
    copying pong.bin from ROMS/Video Olympics - Pong Sports (Paddle) (1977) (Atari, Joe Decuir - Sears) (CX2621 - 99806, 6-99806, 49-75104) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/pong.bin
    copying wizard_of_wor.bin from ROMS/Wizard of Wor (1982) (CBS Electronics - Roklan, Joe Hellesen, Joe Wagner) (M8774, M8794) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/wizard_of_wor.bin
    copying yars_revenge.bin from ROMS/Yars' Revenge (Time Freeze) (1982) (Atari, Howard Scott Warshaw - Sears) (CX2655 - 49-75167) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/yars_revenge.bin
    copying zaxxon.bin from ROMS/Zaxxon (1983) (Coleco) (2454) ~.bin to /usr/local/lib/python3.7/dist-packages/atari_py/atari_roms/zaxxon.bin

## Packages
First, let’s import needed packages. Firstly, we need game (gym)
neural networks (torch.nn)
optimization (torch.optim)

```
import gym
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from collections import deque
import numpy as np
import random
import datetime
```

```
reward_number = 50
```

DQN Algorithm
-------------
We train a policy that tries to maximize the discounted,
cumulative reward
$R_{t_0} = \sum_{t=t_0}^{\infty} \gamma^{t - t_0} r_t$, where
$R_{t_0}$ is *return*. The discount, $\gamma$ is the discount, between $0$ and $1$

Q-learning tries to find a function
$Q^*: State \times Action \rightarrow \mathbb{R}$, maximizes rewards:

\begin{align}\pi^*(s) = \arg\!\max_a \ Q^*(s, a)\end{align}

However, we don't know $Q^*$. So, we use neural network as a approximators, we can simply create one and train it to resemble $Q^*$.

For our training update rule, we'll use a fact that every $Q$
function for some policy obeys the Bellman equation:

\begin{align}Q^{\pi}(s, a) = r + \gamma Q^{\pi}(s', \pi(s'))\end{align}

The difference between the two sides of the equality is known as the temporal difference error, $\delta$:

\begin{align}\delta = Q(s, a) - (r + \gamma \max_a Q(s', a))\end{align}

```
class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.convolution1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3)
        self.convolution2 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5)
        self.convolution3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=7)
        self.fc1 = nn.Linear(in_features=1792, out_features=256)
        self.fc2 = nn.Linear(in_features=256, out_features=128)
        self.fc3 = nn.Linear(in_features=128, out_features=256)
        self.fc4 = nn.Linear(in_features=256, out_features=32)
        self.fc5 = nn.Linear(in_features=32, out_features=9)

    def forward(self, x):
        x = x.cuda()
        x = F.relu(F.max_pool2d(self.convolution1(x), 3))
        x = F.relu(F.max_pool2d(self.convolution2(x), 3))
        x = F.relu(F.max_pool2d(self.convolution3(x), 3, 2))
        x = x.reshape(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.relu(self.fc4(x))
        x = self.fc5(x)
        return x

        ## Quick exercise: Please draw the CNN arch.
```

```
model = CNN()
```

##function to ensure that our code uses the GPU if available, and defaults to using the CPU if it isn't.

```
#function to ensure that our code uses the GPU if available, and defaults to using the CPU if it isn't.
def get_default_device():
    """Pick GPU if available, else CPU"""
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')

# a function that can move data and model to a chosen device.
def to_device(data, device):
    """Move tensor(s) to chosen device"""
    if isinstance(data, (list,tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)
```

## Loss Function and Optimizer

```
criterion = ## your loss function from nn
optimizer = ## your optimizer
```

## DQN Agent
### Replay Memory
We’ll be using experience replay memory for training our DQN. It stores the transitions that the agent observes, allowing us to reuse this data later. By sampling from it randomly, the transitions that build up a batch are decorrelated. It has been shown that this greatly stabilizes and improves the DQN training procedure.
### Hyperparameters

```
class DQNAgent:
    def __init__(self, action_size = 9):
        self.state_size = 3
        self.action_size = action_size
        self.memory_n = deque(maxlen=2000)
        self.memory_p = deque(maxlen=2000)
        self.gamma = 0.95    # discount rate
        self.epsilon = 1.0  # exploration rate
        self.epsilon_min = 0.2
        self.epsilon_decay = 0.995
        self.learning_rate = 0.001
        self.model = model

    def remember(self, state, action, reward, next_state, done):
        if reward == 0:
            self.memory_p.append((state, action, reward, next_state, done))
        else:
            self.memory_n.append((state, action, reward, next_state, done))
    def act(self, state):
        if np.random.rand() <= self.epsilon:
            return random.randrange(self.action_size)
        state_tensor = torch.from_numpy(state).float()
        act_values = self.model(state_tensor).cpu().detach().numpy()
        return np.argmax(act_values[0])  # returns action

    def replay(self, batch_size):
        if len(agent.memory_n) > batch_size / 2:
            print("Negative batch ready:")
            minibatch_n = random.sample(self.memory_n, 5)
            minibatch_p = random.sample(self.memory_p, 59)
            minibatch = random.sample((minibatch_p+minibatch_n), batch_size)
        else:
            minibatch = random.sample(self.memory_p, batch_size)
        for state, action, reward, next_state, done in minibatch:
            ns_model = self.model(torch.from_numpy(next_state).float()).cpu().detach().numpy()
            if reward == 0:
                reward = 1.0001
                target = reward * np.amax(ns_model[0])
                target_f = ns_model
                target_f[0][np.argmax(ns_model[0])] = target
            else:
                reward = reward_number
                target = reward * np.amin(ns_model[0])
                target_max = 0.0001 * np.amax(ns_model[0])
                target_f = ns_model
                target_f[0][action] = target
                target_f[0][random.choice([i for i in range(0,9) if i not in [action]])] = target_max
                self.train(next_state, target_f, epochs=1)
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

    def train(self, input, target, epochs = 1):
        input = torch.from_numpy(input).float().cuda()
        target = torch.from_numpy(target).float().cuda()
        y_pred = 0
        for t in range(1):
            y_pred = model(input)
            loss = - criterion(y_pred, target)
            # print(t, loss.item())
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    def load_all(self, name):
        loaded = torch.load(name)
        self.memory_n = loaded['memory_n']
        self.memory_p = loaded['memory_p']
        self.model.load_state_dict(loaded['state'])

    def save_all(self, name):
        torch.save({'state': self.model.state_dict(),
                    'memory_n': self.memory_n,
                    'memory_p': self.memory_p
                   }, name)

    def load(self, name):
        self.model.load_state_dict(torch.load(name))

    def save(self, name):
        torch.save(self.model.state_dict(), name)

```

```
env = gym.make('MsPacman-v0')
state_size = env.observation_space.shape
action_size = env.action_space.n

```

```
agent = DQNAgent()
```

```
done = False
batch_size = 64

```

## Main training loop

```
EPISODES = 100
for e in range(EPISODES):
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    vw = cv2.VideoWriter('gdrive/' + "Reward_number_" + str(reward_number) + "_" + str(e) + str(datetime.datetime.now()) +  '.avi', fourcc, 4, (160,210))
    state = env.reset()
    state = np.reshape(state, (1, 210,160,3)).transpose(0,3,1,2)/255
    for time in range(1000000000):
        print(time)
        action = agent.act(state)
        next_state, reward, done, _ = env.step(action)
        vw.write(next_state)
        reward = reward if not done else 10
        reward = reward if reward ==0 else 10
        if reward != 0:
            print("ATTENTION NEGATIVE REWARD",reward)
        next_state = np.reshape(next_state, (1, 210,160,3)).transpose(0,3,1,2)/255
        agent.remember(state, action, reward, next_state, done)
        state = next_state
        if done:
            vw.release()
            agent.save('gdrive' + "Reward_number_" + str(reward_number) + "_""Frames:_" + str(time) + "_Episode_" +str(e) + "_Date_" + str(datetime.datetime.now()) + '.pt')
            print("episode: {}/{}, score: {}, e: {:.2}".format(e, EPISODES, time, agent.epsilon))
            break
        if (len(agent.memory_p) > batch_size) & (len(agent.memory_n) > batch_size/2) :
            agent.replay(batch_size)
```


[](https://github.com/liushiru/End-to-end-ASR-Pytorch)

# Project Members

### Module Code: **CS5242**

#### Semester: AY2021-22 Sem 1

#### Group 40

- Liu Shiru (A0187939A)
- Lim Yu Rong, Samuel (A0183921A)
- Yee Xun Wei (A0228597L)

# **README: Important Notes**

This project can **NOT** be trained or run in this notebook. This is because it was forked from a project repository and has a fair amount of setup. Refer to [our Github fork](https://github.com/liushiru/End-to-end-ASR-Pytorch "Our fork") for the full project.

---

To run this project:
1. Clone the repository.
2. Install requirements via `~/requirements.txt`.

## Dataset

Data is collected from the [Ted2srt webpage](https://ted2srt.org "Ted2srt Homepage").

Run `python3 scraper/preprocess.py` from root direcory to scrape and generate dataset.
The script will:
1. Scrape data from website.
2. Preprocess the data
3. Split the data into train-dev-test sets.

Scraped data is saved at `scraper/data/`, processed data will be saved to `data/`.
Alternatively, download data [here](https://nusu-my.sharepoint.com/:u:/g/personal/e0674581_u_nus_edu/EdCosjueRfVIm2kQK3RKBdYBrtlVWU8hsBzcjG7cb9Uo3w?e=fkWh09).

## Training

To train each model:

1. In the root directory, run the command `python3 main.py --config config/<dataset>/<config_file>.yaml --njobs 8`.

### Configuration Files

To use our dataset, set `<dataset>` as `ted` to use scraped data, or `libri` to use public data from [OpenSRL](https://www.openslr.org/12/).

Configuration files are stored as:

#### Train Extractors

| Extractor | Classifier | Configuration file |
|-----------|------------|--------------------|
| MLP       | RNN        | mlp_rnn.yaml       |
| CNN       | RNN        | cnn_rnn.yaml       |
| ANN       | RNN        | ann_rnn.yaml       |
| RNN       | RNN        | rnn_rnn.yaml       |

#### Train Classifiers

| Extractor | Classifier | Configuration file |
|-----------|------------|--------------------|
| CNN       | MLP        | cnn_mlp.yaml       |
| CNN       | CNN        | cnn_cnn.yaml       |
| CNN       | ANN        | cnn_ann.yaml       |

**We advise that a Linux machine should be used**, since some elements in the project may be buggy when used with Windows or even MacOS.

Further information can be found in `~/README.md`.

---

# Problem Statement and Motivation

The Internet of Things (IoT) is becoming relevant in daily life as a natural interface for literate and illiterate users. One important feature of IoT is its ability to provide a hands-free solution for navigating daily tasks: e.g., being able to drive and simultaneously give commands to an IoT device. Speech recognition technology is therefore a key step in the field of IoT.

**Automatic Speech Recognition (ASR)** is traditionally done through a Gaussian Mixture Model-Hidden Markov Model process. However, the rise of Neural Networks (NNs) has completely changed the field of ASR. Most modern ASR software is built on a combination of different NN structures. For this project, we therefore aim to study the performance of individual end-to-end NN structures (e.g. MLP, RNN) independently in this task, to evaluate their strengths and compare them.

Specifically, we aim to test the effectiveness of 4 types of NN architectures (MLP, CNN, RNN, ANN) in correctly parsing audio sentences of variable length. An expected input would be an audio recording of “This is a test sentence”, and the corresponding expected output would be the text “This is a test sentence”.

The codebase used in this report is forked from the [End-to-end-ASR repository](https://github.com/Alexander-H-Liu/End-to-end-ASR-Pytorch "End-to-end-ASR GitHub") by Alexander Liu.

# Data Collection

### Data Source

Data was collected from the [Ted2srt webpage](https://ted2srt.org "Ted2srt Homepage").

Each talk is originally presented as a video, with a single speaker talking. The videos are typically several minutes long.

Ted2srt provides download links for the video, as well as timestamp-annotated transcripts (i.e. SRTs) for most of its videos. These can be seen in the sidebar of [this sample talk](https://ted2srt.org/talks/anil_seth_how_your_brain_invents_your_self "Sample TED Talk"), for example.

---

### Scraper Inputs/Outputs

The two outputs of the scraper are:
- Audio file containing speech
- SRT text file containing timestamp-annotated transcript

Since no audio download link is provided, the scraper scrapes the lowest-quality video, and subsequently converts it into an audio file. The video file is deleted, to save storage space. See the function `download_and_save_video_file()`.

The SRT transcript can be scraped normally. See the function `download_and_save_srt_transcript_text()`. Refer to the 'Data Preprocessing' section for more details on its format.

Any talk which does not contain either the video or the SRT will not be scraped.

---

### Scraper code description

`generate_all_srt_video_urls()` takes advantage of a URL-generation hack to generate any number of ted2srt URLs.

`scrape_data_from_srt_url()`, `download_and_save_video_file()`, and `download_and_save_srt_transcript_text()` all use **BeautifulSoup** to parse the webpage, and the **requests** module to get data from the webpage.

The files downloaded are saved to a `data` folder and named after the title of the TED Talk.

The code for the scraper is shown below.

```python
from bs4 import BeautifulSoup
from moviepy import editor
import requests
from urllib import request, error

from ast import literal_eval
import logging
import os
import shutil
import transcript

os.chdir(os.path.dirname(__file__))

# Logger setup
logger = logging.getLogger()
logging.basicConfig(level="INFO", format="%(levelname)s: %(filename)s: %(message)s")

# Storage locations
DATA_STORAGE_ROOT = os.path.join(os.getcwd(), "data")

# Scraper constants
TED_SRT_HOMEPAGE = "https://ted2srt.org/"
TED_SRT_TALKS = "talks/"
TED_SRT_TALKPAGE = TED_SRT_HOMEPAGE + TED_SRT_TALKS

def main(number_of_talks=200, starting_video_id=100):
    if not os.path.isdir(DATA_STORAGE_ROOT):
        try:
            os.makedirs(DATA_STORAGE_ROOT)
        except FileNotFoundError:
            logging.error("Unable to create data storage folder.")
            raise Exception("Ensure that the current working directory is set to the project root.")

    _number_of_talks = number_of_talks  # Number of talks to attempt to scrape from
    _starting_video_id = starting_video_id  # Start retrieving from this video id

    logging.info("Parsing audio files with SRT transcripts from TED2SRT website...")
    video_urls = generate_all_srt_video_urls(TED_SRT_TALKPAGE, _number_of_talks, starting_video_id=_starting_video_id)
    for url in video_urls:
        scrape_data_from_srt_url(url)

## Code that handles scraping from the TED2SRT website

def generate_all_srt_video_urls(base_url, number_of_urls, starting_video_id=1):
    """
    Generates a list of urls to attempt to scrape from. This is designed to work specifically for the TED2SRT page.
    This method will not check whether the links are valid. It will only generate the required links.

    This method works because the TED2SRT page also stores videos mapped to their id. For example, visiting the webpage
    "homepage/talks/1" returns the 1st video in id.

    :param base_url: The base url of the TED2SRT talks page.
    :param number_of_urls: The number of urls to generate.
    :param starting_video_id: The starting video id, which defaults to the 1st video.
    :return: Returns a list of urls.
    """
    logging.info(f"Generating {number_of_urls} starting from video id {starting_video_id}...")
    logging.warning("Many of these links might not work. Consider changing the value of _number_of_talks in line 52.")
    return [base_url + str(video_id) for video_id in range(starting_video_id, starting_video_id + number_of_urls)]

# TED2SRT specific page constants
VIDEO_NAME_KEY = "slug"

VIDEO_DOWNLOAD_SUB_URL_KEY = "mediaSlug"
START_OF_AUDIO_DOWNLOAD_URL = "https://download.ted.com/talks/"
END_OF_AUDIO_DOWNLOAD_URL = "-320k.mp4"

VIDEO_SRT_ID_SUB_URL_KEY = "id"
START_OF_SRT_TRANSCRIPT_URL = "https://ted2srt.org/api/talks/"
END_OF_SRT_TRANSCRIPT_URL = "/transcripts/download/srt?lang=en"

def scrape_data_from_srt_url(url):
    """
    Attempts to scrape the audio file and SRT transcript from the given webpage.
    Specifically written to scrape data from the TED2SRT page (root: https://ted2srt.org/).

    If either the audio file or the SRT transcript does not exist, then this method will not scrape anything.
    Otherwise, the data will be saved to a folder marked by the url.
    :param url: The url to scrape data from.
    :return: No return value.
    """
    logger.info(f"\n\nScraping from url: <{url}>")
    response = requests.get(url)
    page_soup = BeautifulSoup(response.text, "html.parser")

    # Identify the metadata object and parse the key phrases needed for video and SRT download
    talk_metadata_script = page_soup.select("script")
    if not talk_metadata_script:
        logging.warning(f"The url {url} is invalid or has no metadata object. Continuing...")
        return

    talk_metadata_string = talk_metadata_script[0].string
    talk_metadata_object = literal_eval(talk_metadata_string.split(" = ")[1])

    srt_download_keyword = str(talk_metadata_object[VIDEO_SRT_ID_SUB_URL_KEY])
    video_download_keyword = str(talk_metadata_object[VIDEO_DOWNLOAD_SUB_URL_KEY])
    name_of_video = str(talk_metadata_object[VIDEO_NAME_KEY])

    srt_download_url = START_OF_SRT_TRANSCRIPT_URL + srt_download_keyword + END_OF_SRT_TRANSCRIPT_URL
    video_download_url = START_OF_AUDIO_DOWNLOAD_URL + video_download_keyword + END_OF_AUDIO_DOWNLOAD_URL

    saved_folder_path = os.path.join(DATA_STORAGE_ROOT, name_of_video)

    success_status = download_and_save_srt_transcript_text(srt_download_url, saved_folder_path)
    if not success_status:
        return

    download_and_save_video_file(video_download_url, saved_folder_path)

# Video constants
VIDEO_FILE_EXTENSION = "video.mp4"
AUDIO_FILE_EXTENSION = "audio.mp3"
SRT_TRANSCRIPT_FILE_EXTENSION = "srt_transcript.txt"

def download_and_save_video_file(video_download_url, path_to_save_to):
    """
    Downloads a video file, given its url link. Converts the video file to audio. Saves the file to the required folder.
    Specifically designed for the TED2SRT website. Note that no audio download exists. Therefore, the video must be
    downloaded first, and then converted into audio. To save space, the video file is then deleted.
    :param video_download_url: The url to download the video from.
    :param path_to_save_to: The folder name to save this file to.
    :return: Returns True if the download was successful, and False otherwise.
    """
    if not os.path.isdir(path_to_save_to):
        os.mkdir(path_to_save_to)

    video_file_saved_location = os.path.join(path_to_save_to, VIDEO_FILE_EXTENSION)
    try:
        _, _ = request.urlretrieve(video_download_url, filename=video_file_saved_location)
        logger.info("Video file saved.")
    except error.HTTPError:
        logger.warning("No video file of lowest quality found. Deleting this specific data folder...")
        srt_transcript_file_saved_location = os.path.join(path_to_save_to, SRT_TRANSCRIPT_FILE_EXTENSION)

        if os.path.isfile(srt_transcript_file_saved_location):
            shutil.rmtree(path_to_save_to)

        if os.path.isdir(path_to_save_to):
            os.rmdir(path_to_save_to)

        return False

    logger.info("Converting video.mp4 file to audio.mp3 file...")
    audio_file_saved_location = os.path.join(path_to_save_to, AUDIO_FILE_EXTENSION)

    video = editor.VideoFileClip(video_file_saved_location)
    audio = video.audio
    audio.write_audiofile(audio_file_saved_location, verbose=False, logger=None)
    video.close()

    os.remove(video_file_saved_location)
    logging.info("Audio file saved. Video file deleted.")
    return True

def download_and_save_srt_transcript_text(srt_transcript_url, path_to_save_to):
    """
    Saves the transcript text to the required folder as a file. This is specifically designed to work with the TED2SRT
    website.
    :param srt_transcript_url: The url to download the srt transcript from.
    :param path_to_save_to: The folder name to save this file to.
    :return: Returns True if the download was successful, and False otherwise.
    """
    if not os.path.isdir(path_to_save_to):
        os.mkdir(path_to_save_to)

    srt_transcript_file_saved_location = os.path.join(path_to_save_to, SRT_TRANSCRIPT_FILE_EXTENSION)
    try:
        _, _ = request.urlretrieve(srt_transcript_url, filename=srt_transcript_file_saved_location)

        # Check for SRT accuracy before downloading audio
        srt_is_accurate = transcript.check_srt_accuracy(srt_transcript_file_saved_location)
        if not srt_is_accurate:
            logger.warning("SRT transcript not accurate. Deleting this specific data folder...")
            shutil.rmtree(path_to_save_to)
            return False

        logger.info("SRT transcript saved.")
        return True
    except error.HTTPError:
        logger.warning("No SRT transcript found. Deleting this specific data folder...")
        audio_file_saved_location = os.path.join(path_to_save_to, VIDEO_FILE_EXTENSION)

        if os.path.isfile(audio_file_saved_location):
            shutil.rmtree(path_to_save_to)

        if os.path.isdir(path_to_save_to):
            os.rmdir(path_to_save_to)

        return False

main()

```

# Data Preprocessing

Preprocess scraped data to input into `Dataset` and `DataLoader`. Includes data cleaning, cutting audio into multiple audio slices according to SRT annotated time and prepare label for each audio slice.

---

### Data Preprocessing Inputs/Outputs

- Input from scraped audio files and SRT.
- Output preprocessed ready dataset for ASR.

---

### Data Preprocessing Descriptions

Symbols are removed from label and converted to lowercase.

Data that are less accurate are removed. Checking done for SRT that starts at the same time (e.g. `00:00:12,820 --> 00:00:14,820`). SRT that does not include introduction music time is filtered. Laughter and applause is removed.

Raw SRT snippet
```
1
00:00:12,312 --> 00:00:14,493
Six months ago, I got an email

2
00:00:14,493 --> 00:00:15,900
from a man in Israel
```
converted to
```
<audio_id>-1 six months ago i got an email
<audio_id>-2 from a man in israel
```
and stored at `<audio_id>.trans.txt`. Corresponding sliced audio files are named `<audio_id>_<audio_index>.mp3`.

After `build_dataset()` has preprocessed the data, the data is split into  train-dev-test sets.

```python
import os
from joblib import Parallel, delayed
from os.path import join
from pathlib import Path
import random
import shutil

import scraper
import remove_applause
import re
import subprocess
import num2words
import pydub
from pydub import AudioSegment
from tqdm import tqdm
import logging
logger = logging.getLogger()
logging.basicConfig(level="INFO", format="%(levelname)s: %(filename)s: %(message)s")

random.seed(42)
AUDIO_EXTENSION = 'mp3'
os.chdir(os.path.dirname(__file__))
corpus_path = '../data/TedSrt'
txt_norm = 'tedsrt-norm.txt'
txt_vocab = 'tedsrt-vocab.txt'
processed_data_path = join(corpus_path, 'processed_data') # store all data in this folder before splitting
txt_norm_path = join(corpus_path, txt_norm)
txt_vocab_path = join(corpus_path, txt_vocab)

src_path = 'data/' # raw data source
ratio=[.6, .2, .2] # train test split

# filter audio and srt
duration_diff = 7 # seconds
repeated_occurrence = 50

def clean_text(text):
    '''
    Text processing to clean text before saving as label
    to lowercase, convert years to words, convert digits to words, remove symbols
    '''
    text = text.lower().strip('\n')
    text = re.sub(r'[^\w\s]', ' ', text)
    text = ' '.join([num2words.num2words(i, to='year') if (i.isdigit() & (len(i) == 4)) else i for i in text.split()]) # year to words
    text = ' '.join([num2words.num2words(i) if i.isdigit() else i for i in text.split()]) # num to words
    text = re.sub(' +', ' ', text) # remove redundant spaces
    text = text.replace('-', ' ')
    return text

def normalize_text(text):
    '''
    Text processing to normalize text for language model training
    should transform word numeral to numeric, normalize date formats, lemmatization; most not implemented currently
    '''
    text = text.lower().strip('\n')
    text = re.sub(r'[^\w\s]', ' ', text)
    text = re.sub(' +', ' ', text) # remove redundant spaces
    text = text.replace('-', ' ')
    return text

def to_ms(string):
    '''
    Convert string '00:00:00,000' to milliseconds
    to be used for audio slicing
    '''
    string = string.replace(',','')
    hour, minute, second = string.split(':')
    second = int(second)
    second += int(hour) * 3600 * 1000
    second += int(minute) * 60 * 1000
    second = second
    return second

def txt_to_trans(txt_file, file_name):
    '''
    Convert txt file to transcript format ready to be read into Dataset
    lines formatted as 'filename-idx text_label'
    return lines and time_slices
    '''
    file = open(txt_file, 'r')
    lines = file.readlines()
    file.close()

    transcript = [] # label for audio
    txt_src = [] # label for language model
    time_slices = []

    for i in range(len(lines)):
        idx = re.search('^[\d]+$', lines[i].strip('\ufeff'))
        if idx:
            idx = idx[0]
            time_frame = re.findall('[0-9]{2}:[0-9]{2}:[0-9]{2},[0-9]{3}', lines[i+1])
            if time_frame:
                start, end = to_ms(time_frame[0]), to_ms(time_frame[1])
                time_slices.append((idx, (start, end)))

                audio_label = lines[i+2]
                audio_label = clean_text(audio_label)
                new_line = f"{file_name}-{idx} {audio_label}"
                transcript.append(new_line)

                lm_label = normalize_text(audio_label)
                txt_src.append(lm_label)

    return transcript, time_slices, txt_src

def save_txt(txt, output_path):
    '''
    Save transcript to output_path
    '''
    if not os.path.exists(os.path.dirname(output_path)):
        os.makedirs(os.path.dirname(output_path))

    with open(output_path, 'a+') as f:
        for line in txt:
            f.write(f"{line}\n")
        f.close()

def check_ms_accuracy(time_slices):
    '''
    Check whether srt is accurate, remove those not accurate to milliseconds, check repeated occurrence of 820
    '''
    for time_slice in time_slices:
        occurrence_820 = sum(1 for elem in time_slices if elem[1][0] % 1000 == 820)
    if occurrence_820 > repeated_occurrence:
        return occurrence_820
    return 0

def check_intro_timing(audio_duration, audio_srt_duration):
    '''
    Check whether srt matches audio, remove those without taking water drop intro into account
    '''
    duration_diff = audio_duration - audio_srt_duration
    if duration_diff > duration_diff:
        return duration_diff
    return 0

def build_dataset(src_path=src_path, sample=100):
    '''
    Build dataset from raw data scraper/data/
    split audio into slices
    convert srt to trans.txt ready for training
    '''
    shutil.rmtree(corpus_path, ignore_errors=True)
    folders = os.listdir(src_path)[:100]
    for idx, curr_folder in enumerate(tqdm(folders, desc="Building dataset from scraped data")):
        file_name = str(idx) #save the transcript as num, can be changed to folder name
        output_path = join(processed_data_path, file_name)
        txt_output_path = join(output_path, file_name + '.trans.txt')

        try:
            txt_path = list(Path(join(src_path, curr_folder)).rglob('*.txt'))[0]
            audio_path = list(Path(join(src_path, curr_folder)).rglob('*.' + AUDIO_EXTENSION))[0]
        except:
            tqdm.write(f'Data not complete {idx} {curr_folder}')
            continue
        transcript, time_slices, txt_src = txt_to_trans(txt_path, file_name)

        audio_file = AudioSegment.from_file(audio_path, AUDIO_EXTENSION)

        # check timing accuracy by milliseconds
        ms_not_accurate = check_ms_accuracy(time_slices)
        if ms_not_accurate:
            shutil.rmtree(output_path, ignore_errors=True)
            continue

        # check timing of audio intro
        audio_duration = audio_file.duration_seconds
        audio_srt_duration = time_slices[-1][-1][-1] / 1000
        intro_not_matched = check_intro_timing(audio_duration, audio_srt_duration)
        if intro_not_matched:
            shutil.rmtree(output_path, ignore_errors=True)
            continue

        # writing output
        save_txt(transcript, txt_output_path)
        save_txt(txt_src, txt_norm_path)
        for idx, time_slice in time_slices:
            audio_slice = audio_file[time_slice[0]:time_slice[1]]
            audio_output_path = join(output_path, f"{file_name}-{idx}.{AUDIO_EXTENSION}")
            audio_slice.export(audio_output_path, format=AUDIO_EXTENSION)

        tqdm.write(f'Successfully created {idx} {curr_folder}')

def move_folders(folders, source, destination):
    '''
    Move a list of folders from source to destination
    '''
    logger.info(f'{destination.split("/")[-1]}: {folders}')
    for folder in folders:
        shutil.move(join(source, folder), join(destination, folder))

def split_dataset(ratio=[.6, .2, .2]):
    '''
    Split dataset to train dev test according to ratio
    '''
    folders = os.listdir(processed_data_path)
    folders.sort()
    random.shuffle(folders)

    split_1 = int(ratio[0] * len(folders))
    split_2 = int((ratio[0] + ratio[1]) * len(folders))

    move_folders(folders[:split_1], processed_data_path, join(corpus_path, 'train'))
    move_folders(folders[split_1:split_2], processed_data_path, join(corpus_path, 'dev'))
    move_folders(folders[split_2:], processed_data_path, join(corpus_path, 'test'))
    shutil.rmtree(processed_data_path, ignore_errors=True)

def main():
    print('Scraping data...')
    scraper.main(number_of_talks=500, starting_video_id=2000) # 1000-1500

    print('\nStart preprocessing...')
    build_dataset(src_path, sample=100)
    split_dataset(ratio=ratio)

    print("\nRemoving laughter and applause audio files...")
    remove_applause.main(corpus_path)

main()
```

## PyTorch Datasets

`TedDataset` maps the sliced audio files and their corresponding labels.

`TedDataset` reads all audio files and search for corresponding labels stored at `<audio_id>.trans.txt`.

```python
from tqdm import tqdm
from pathlib import Path
from os.path import join, getsize
from ntpath import basename
from joblib import Parallel, delayed
from torch.utils.data import Dataset

# Additional (official) text src provided
OFFICIAL_TXT_SRC = ['tedsrt-lm-norm.txt']
# Default num. of threads used for loading TedSrt
READ_FILE_THREADS = 1
# Default audio extension
AUDIO_EXTENSION = 'mp3'

def read_text(file):
    '''Get transcription of target mp3 file,
       it's somewhat redundant for accessing each txt multiplt times,
       but it works fine with multi-thread'''
    src_file = '-'.join(file.split('-')[:-1])+'.trans.txt'
    idx = basename(file).split('.')[0]

    with open(src_file, 'r') as fp:
        for line in fp:
            if idx == line.split(' ')[0]:
                return line[:-1].split(' ', 1)[1]

class TedDataset(Dataset):
    def __init__(self, path, split, tokenizer, bucket_size, ascending=False):
        # Setup
        self.path = path
        self.bucket_size = bucket_size

        # List all wave files
        file_list = []
        for s in split:
            split_list = list(Path(join(path, s)).rglob("*." + AUDIO_EXTENSION))
            assert len(split_list) > 0, "No data found @ {}".format(join(path,s))
            file_list += split_list
        # Read text
        text = Parallel(n_jobs=READ_FILE_THREADS)(
            delayed(read_text)(str(f)) for f in file_list)
        text = [tokenizer.encode(txt.upper()) for txt in text]

        # Sort dataset by text length
        self.file_list, self.text = zip(*[(f_name, txt)
                                          for f_name, txt in sorted(zip(file_list, text), reverse=not ascending, key=lambda x:len(x[1]))])

    def __getitem__(self, index):
        if self.bucket_size > 1:
            # Return a bucket
            index = min(len(self.file_list)-self.bucket_size, index)
            return [(f_path, txt) for f_path, txt in
                    zip(self.file_list[index:index+self.bucket_size], self.text[index:index+self.bucket_size])]
        else:
            return self.file_list[index], self.text[index]

    def __len__(self):
        return len(self.file_list)
```

# Data Exploration

#### Audio metadata of an audio slice
```
Sample Rate: 44100
Shape: (1, 84055)
Dtype: torch.float32
 - Max:          0.523
 - Min:         -0.319
 - Mean:        -0.000
 - Std Dev:      0.081
 ```

#### Audio signal

Waveform plot of sample audio signal with length of 1.9s. Duration length can obtained: `signal_frames / sample_rate`.

#### Spectrogram of raw audio signal

# Data Processing

Steps to compute filter banks are motivated to mimic how human perceives audio signals<sup>[[1]](#reference)</sup>.

1. Apply pre-emphasis filter on audio signal (amplify the high frequencies since high frequencies have smaller magnitude).
2. Cut signal into window frames (assume signal is stationary over a short period time).
3. Compute the power spectrum of the signal (Fourier transform) for each window.

Kaldi filter banks transformation applied on audio signals. 40 mel coefficients is kept.

```python
feat_dim = 40
waveform_trans = torchaudio.compliance.kaldi.fbank(signal, frame_length=25, frame_shift=10, num_mel_bins=feat_dim)
plot_spectrogram(waveform_trans.transpose(0, 1).detach(), title="Filter Banks", ylabel='mel bins')
```

# Model Design

## Description of End-to-end-ASR model

In this project, we are training a feature extractor and classifier, that collectively classify words based on audio inputs. The feature extractor mechanism extracts features from timesteps (i.e. speech frames). This feature vector is used by the classifier to generate an output sequence of characters. The output sequence is evaluated using Connectionist Temporal Classification (CTC) Loss and Word Error Rate (WER).

### Baseline model

The combination of extractor and classifier used in the ASR-End-to-end paper is a VGGNet followed by an LSTM. We consider VGGNet to be a CNN, and LSTM to be an RNN.

The use of a CNN feature extractor is motivated by previous studies<sup>[[2]](#reference)[[3]](#reference)</sup>, and similarly for the use of an LSTM classifier<sup>[[4]](#reference)</sup>.

## Experimentation Plan

The following table describes the 7 sets of extractor-classifier combinations used in the tests. The first combination (Baseline) is used by the ASR-End-to-end paper. Apart from that, there are 6 other combinations that we can use, that involve changing either the extractor or the classifier of the original CNN-RNN combination, but not both.

##### Baseline

Model Tested | Extractor | Classifier
--- | --- | ---
Baseline | CNN | RNN

##### Change the Extractor

Model Tested | Extractor | Classifier
--- | --- | ---
MLPExtractor | MLP | RNN
RNNExtractor | RNN | RNN
ANNExtractor | ANN | RNN

##### Change the Classifier

Model Tested | Extractor | Classifier
--- | --- | ---
MLPClassifier | CNN | MLP
CNNClassifier | CNN | CNN
ANNClassifier | CNN | ANN

We cannot change both the extractor and classifier in a single test because we cannot compare it easily with the original baseline.

## Code Written

```python
# TAKE NOTE that these code elements cannot be trained within this notebook, due to the complex project ecosystem.
# Refer to the 'Important Notes' section at the start of this notebook for instructions on running the project.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init
```

### Extractors

Extractor generates a sequence of feature vectors. Each feature vector is extracted from a small overlapped window of audio frames. Extractor transforms $x$, to high-level representation $h = (h_1, ..., h_L)$.

Extractor includes downsampling of timesteps.

For example downsampling by a factor of 4 from 523 timesteps to 130 timesteps in RNN extractor. Downsampling is also achieved by MaxPooling of CNN extractors.

#### MLP Extractor

```python
class MLPExtractor(nn.Module):
    ''' A simple MLP extractor for acoustic feature down-sampling
    Every 4 frame will be convert to corresponding features by MLP'''

    def __init__(self, input_dim, out_dim):
        super(MLPExtractor, self).__init__()

        self.out_dim = out_dim
        self.hide_dim = input_dim * 3
        self.extractor = nn.Sequential(
            nn.Linear(input_dim, self.hide_dim),
            nn.ReLU(),
            nn.Linear(self.hide_dim, self.hide_dim),
            nn.ReLU(),
            nn.Linear(self.hide_dim, self.hide_dim),
            nn.ReLU(),
            nn.Linear(self.hide_dim, self.hide_dim),
            nn.ReLU(),
            nn.Linear(self.hide_dim, self.out_dim),
        )

    def reshape_input(self, feature, group_size):
        down_sample_len = feature.size(1) // group_size
        feature = feature[:,:down_sample_len*group_size,:]
        reshape_feature = feature.reshape(feature.size(0) * down_sample_len, group_size*feature.size(2))
        return reshape_feature

    def forward(self, feature, feat_len):
        bs = feature.size(0)
        feature = self.reshape_input(feature, group_size=4)
        raw_output = self.extractor(feature)

        reshape_output = raw_output.reshape(bs, raw_output.size(0) // bs, self.out_dim)
        return reshape_output, torch.div(feat_len,4, rounding_mode="floor")
```

#### RNN Extractor

Applied RNN layer on input (speech frames/ timesteps). RNN output is followed by downsampling by factor of 4.

```python
class RNNExtractor(nn.Module):
    ''' A 2-layer RNN extractor for acoustic feature down-sampling'''

    def __init__(self, input_dim, out_dim):
        super(RNNExtractor, self).__init__()

        self.out_dim = out_dim
        self.layer = nn.RNN(input_dim, self.out_dim)
        self.dp = nn.Dropout(p=0.2)
        self.linear = nn.Linear(self.out_dim, self.out_dim)

    def forward(self, feature, feat_len):

        # feat_len = feat_len//4
        feat_len = torch.div(feat_len, 4, rounding_mode='floor')
        feature, _ = self.layer(feature)
        # Normalization
        feature = self.dp(feature)

        # Downsample timestep
        sample_rate = 4
        if feature.shape[1] % sample_rate != 0:
            feature = feature[:, :-(feature.shape[1] % 4), :].contiguous()
        feature = feature[:, ::sample_rate, :].contiguous()

        return feature, feat_len
```

#### ANN Extractor

In this work we applied local self-attention layer<sup>[[5]](#reference)</sup> to replace spatial convolutions as a feature extractor. The attention is computed as follows

\begin{align}
y_{ij} =
  \sum_{a, b \in \, \mathcal{N}_k(i, j)}
    \texttt{softmax}_{a b}\left(  q_{i j}^\top k_{a b}  \right)  v_{a b}
\end{align}

where the *queries* $q_{ij} = W_Q x_{ij}$, *keys* $k_{ab} = W_K x_{ab}$, and *values* $v_{ab} = W_V x_{ab}$ are linear transformations of the pixel in position $ij$ and the neighborhood pixels<sup>[[5]](#reference)</sup>.

```python
class ANNExtractor(nn.Module):
    ''' ANN extractor with self-attention mechanism from https://github.com/leaderj1001/Stand-Alone-Self-Attention'''

    def __init__(self, input_dim, groups=8):
        super(ANNExtractor, self).__init__()
        self.init_dim = 64
        self.hide_dim = 128
        in_channel, freq_dim, out_dim = self.check_dim(input_dim)
        self.in_channel = in_channel
        self.freq_dim = freq_dim
        self.out_dim = 640
        width = freq_dim

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channel, width, 3, stride=1, padding=1),
            nn.BatchNorm2d(width),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2),  # Half-time dimension
            nn.Conv2d(width, width, 3, stride=1, padding=1),
            nn.MaxPool2d(2, stride=2),  # Half-time dimension
        )
        self.conv2 = nn.Sequential(
            AttentionConv(width, width, kernel_size=7, padding=3, groups=groups),
            nn.BatchNorm2d(width),
            nn.ReLU(),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(width, 64, 3, stride=1, padding=1),
            nn.BatchNorm2d(64),
        )

    def check_dim(self, input_dim):
        # Check input dimension, delta feature should be stack over channel.
        if input_dim % 13 == 0:
            # MFCC feature
            return int(input_dim/13), 13, (13//4)*self.hide_dim
        elif input_dim % 40 == 0:
            # Fbank feature
            return int(input_dim/40), 40, (40//4)*self.hide_dim
        else:
            raise ValueError(
            'Acoustic feature dimension for VGG should be 13/26/39(MFCC) or 40/80/120(Fbank) but got '+input_dim)

    def view_input(self, feature, feat_len):
        # downsample time
        # feat_len = feat_len//4
        feat_len = torch.div(feat_len, 4, rounding_mode='floor')
        # crop sequence s.t. t%4==0
        if feature.shape[1] % 4 != 0:
            feature = feature[:, :-(feature.shape[1] % 4), :].contiguous()
        bs, ts, ds = feature.shape
        # stack feature according to result of check_dim
        feature = feature.view(bs, ts, self.in_channel, self.freq_dim)
        feature = feature.transpose(1, 2)

        return feature, feat_len

    def forward(self, feature, feat_len):

        # Feature shape BSxTxD -> BSxCH(1)xT/4xD
        feature, feat_len = self.view_input(feature, feat_len) #downsample on time
        # Foward
        # BS x 1 x T/4 x D/4 -> BS x width x T/4 x D
        feature = self.conv1(feature)
        # BS x width x T/4 x D/4 -> BS x width x T/8 x D/2 (attention)
        feature = self.conv2(feature)

        # BS x width x T/8 x D/2 -> BS x 64 x T/8 x D/4
        feature = self.conv3(feature)
        feature = F.relu(feature)
        # BS x 64 x T/8 x D/4 -> BS x T/8 x 64 x 8D
        feature = feature.transpose(1, 2)
        # BS x T/8 x 64 x 8D -> BS x T/8 x 8D
        feature = feature.contiguous().view(feature.shape[0], feature.shape[1], -1)

        return feature, feat_len

class AttentionConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, groups=1, bias=False):
        super(AttentionConv, self).__init__()
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.groups = groups

        assert self.out_channels % self.groups == 0, "out_channels should be divided by groups. (example: out_channels: 40, groups: 4)"

        self.rel_h = nn.Parameter(torch.randn(out_channels // 2, 1, 1, kernel_size, 1), requires_grad=True)
        self.rel_w = nn.Parameter(torch.randn(out_channels // 2, 1, 1, 1, kernel_size), requires_grad=True)

        self.key_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)
        self.query_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)
        self.value_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)

        self.reset_parameters()

    def forward(self, x):
        batch, channels, height, width = x.size()

        padded_x = F.pad(x, [self.padding, self.padding, self.padding, self.padding])
        q_out = self.query_conv(x)
        k_out = self.key_conv(padded_x)
        v_out = self.value_conv(padded_x)

        k_out = k_out.unfold(2, self.kernel_size, self.stride).unfold(3, self.kernel_size, self.stride)
        v_out = v_out.unfold(2, self.kernel_size, self.stride).unfold(3, self.kernel_size, self.stride)

        k_out_h, k_out_w = k_out.split(self.out_channels // 2, dim=1)
        k_out = torch.cat((k_out_h + self.rel_h, k_out_w + self.rel_w), dim=1)

        k_out = k_out.contiguous().view(batch, self.groups, self.out_channels // self.groups, height, width, -1)
        v_out = v_out.contiguous().view(batch, self.groups, self.out_channels // self.groups, height, width, -1)

        q_out = q_out.view(batch, self.groups, self.out_channels // self.groups, height, width, 1)

        out = q_out * k_out
        out = F.softmax(out, dim=-1)
        out = torch.einsum('bnchwk,bnchwk -> bnchw', out, v_out).view(batch, -1, height, width)

        return out

    def reset_parameters(self):
        init.kaiming_normal_(self.key_conv.weight, mode='fan_out', nonlinearity='relu')
        init.kaiming_normal_(self.value_conv.weight, mode='fan_out', nonlinearity='relu')
        init.kaiming_normal_(self.query_conv.weight, mode='fan_out', nonlinearity='relu')

        init.normal_(self.rel_h, 0, 1)
        init.normal_(self.rel_w, 0, 1)
```

### Classifiers

Classifier generates an output sequence $(y_1, . . . , y_T)$ from input $h$. $h$ is the output of the extractor. The classifier's output $y$ is a sequence of word tokens. The sequence $y$ is expected to include a dimension of the same size as $h$.

#### MLP Classifier

The MLPClassifier class uses 5 linear layers. For every timeframe, the extractor map the audio fbanks feature to a latent representation. Then the MLP Classifier aims to map the latent representation to a word token prediction.

Unlike original BiLSTM classifier that looks after timeframe from before and after, the MLP classifier is completely local. It only cares about the current timeframe.

```python
class MLPCLassfier(nn.Module):
    def __init__(self, input_dim):
        super(MLPCLassfier, self).__init__()
        h1 = input_dim*2
        h2 = h1*2
        self.out_dim = h2*2
        self.classifier = nn.Sequential(
            nn.Linear(input_dim, h1),
            nn.ReLU(),
            nn.Linear(h1, h2),
            nn.ReLU(),
            nn.Linear(h2, h2),
            nn.ReLU(),
            nn.Linear(h2, h2),
            nn.ReLU(),
            nn.Linear(h2, self.out_dim),
        )

    def reshape_input(self, feature, group_size):
        down_sample_len = feature.size(1) // group_size
        feature = feature[:,:down_sample_len*group_size,:]
        reshape_feature = feature.reshape(feature.size(0) * down_sample_len, group_size*feature.size(2))
        return reshape_feature

    def forward(self, feature):

        return self.classifier(feature)

```

#### CNN Classifier

The CNNClassifer class uses multiple convolutional layers and only one dense layer. The choice of a single dense layer was used to isolate the effectiveness of convolutional layers specifically.

The dense layer cannot be directly connected to the convolutional layers. This is because the input data is of size `(timesteps x features)`, where `features` is fixed at 1280, but `timesteps` differs between inputs. Additionally, the output vector is required to contain a dimension with size equal to `timesteps`.

A workaround for this is to apply the linear layer to each timestep separately, and output a separate value for each of the timesteps. Additionally, pooling layers should not pool along the timestep axis, since the output vector must have a dimension with size equal to `timesteps`.

```python
class CNNClassifier(nn.Module):
    """4-layer CNN classifier module."""

    def __init__(self, input_dim):
        super(CNNClassifier, self).__init__()
        self.hidden_dim_1 = hidden_dim_1 = input_dim // 64
        self.hidden_dim_2 = hidden_dim_2 = hidden_dim_1 * 2

        self.out_dim = input_dim // 4

        conv_hyperparams = {
            "kernel_size": (3,3),
            "dilation": (1,1),
            "stride": (1,1),
            "padding": (1,1),
        }

        self.conv1 = nn.Conv2d(in_channels=1, out_channels=hidden_dim_1, **conv_hyperparams)
        self.conv2 = nn.Conv2d(in_channels=hidden_dim_1, out_channels=hidden_dim_1, **conv_hyperparams)
        self.conv3 = nn.Conv2d(in_channels=hidden_dim_1, out_channels=hidden_dim_2, **conv_hyperparams)
        self.conv4 = nn.Conv2d(in_channels=hidden_dim_2, out_channels=hidden_dim_2, **conv_hyperparams)
        self.dense = nn.Linear(hidden_dim_2 * self.out_dim, self.out_dim)
        self.pool = nn.MaxPool2d((1, 2))

    def reshape_input(self, feature, group_size):
        down_sample_len = feature.size(1) // group_size
        feature = feature[:,:down_sample_len*group_size,:]
        reshape_feature = feature.reshape(feature.size(0) * down_sample_len, group_size*feature.size(2))
        return reshape_feature

    def forward(self, feature):
        # Input size is varied - size is N timesteps x 1280 features. Batch size = 1 and number of input channels = 1.
        # Therefore, inputs must be unsqueezed.
        feature = feature.unsqueeze(dim=0).unsqueeze(dim=1)

        # Format: input_channels x timesteps x features --> output_channels x new_timesteps x new_features
        feature = self.conv1(feature)  # 1 x N x 1280 --> 24 x N x 1280
        feature = F.relu(feature)
        feature = self.conv2(feature)  # 24 x N x 1280 --> 24 x N x 1280
        feature = F.relu(feature)
        feature = self.pool(feature)  # 24 x N x 1280 --> 24 x N x 640
        feature = self.conv3(feature)  # 24 x N x 640 --> 48 x N x 640
        feature = F.relu(feature)
        feature = self.conv4(feature)  # 48 x N x 640 --> 48 x N x 640
        feature = F.relu(feature)
        feature = self.pool(feature)  # 48 x N x 640 --> 48 x N x 320

        feature = feature.transpose(0, 1)  # 48 x N x 320 --> N x 48 x 320
        feature = feature.view(-1, self.hidden_dim_2 * self.out_dim) # N x 48 x 320 --> N x (48 x 320)
        feature = torch.stack([
            self.dense(feature[index].unsqueeze(dim=0))
            for index in range(feature.size()[0])])  # N x (48 x 320) --> N x 320
        feature = feature.transpose(0, 1)  # N x 320 --> 320 x N
        return feature
```

#### ANN Classifier

The ANNClassifier class uses a single Self-Attention layer (referencing [this implementation](https://github.com/leaderj1001/Stand-Alone-Self-Attention/blob/master/attention.py "Self-Attention layer implementation") of a Self-Attention layer). The implementation is identical to the one used in the ANNExtractor.

The ANNClassifier only uses a single dense layer. This is to isolate the effectiveness of the Self-Attention layer.

However, because the Self-Attention layer is extremely space-intensive, only a single Self-Attention layer can be used, and pooling is applied prior to the data being forwarded to it.

Additionally, the ANNClassifier's dense layer faces the same issue as the CNNClassifier - the input data is of unknown size, and therefore the dense layer is applied to each timestep individually.

```python
class ANNClassifier(nn.Module):
    """ANN classifier with Self-Attention module from https://github.com/leaderj1001/Stand-Alone-Self-Attention/blob/master/attention.py."""

    def __init__(self, input_dim):
        super(ANNClassifier, self).__init__()
        self._pool_value = 4
        self.hidden_dim_1 = hidden_dim_1 = 20
        self.linear_input_dim = hidden_dim_1 * input_dim // self._pool_value
        self.out_dim = 80

        attn_hyperparams = {
            "kernel_size": 3,
            "padding": 1,
        }

        self.attn = AttentionConv(1, hidden_dim_1, **attn_hyperparams)
        self.dense = nn.Linear(self.linear_input_dim, self.out_dim)
        self.pool = nn.MaxPool2d((1, self._pool_value))

    def reshape_input(self, feature, group_size):
        down_sample_len = feature.size(1) // group_size
        feature = feature[:,:down_sample_len*group_size,:]
        reshape_feature = feature.reshape(feature.size(0) * down_sample_len, group_size*feature.size(2))
        return reshape_feature

    def forward(self, feature):
        # Input size is varied - size is N timesteps x 1280 features. Batch size = 1 and number of input channels = 1.
        # Therefore, inputs must be unsqueezed.
        feature = feature.unsqueeze(dim=0).unsqueeze(dim=1)

        # Format: output_channels x new_timesteps x new_features

                                                   # 1 x N x 1280
        feature = self.pool(feature)               # 1 x N x 320
        feature = self.attn(feature)               # 20 x N x 320
        feature = F.relu(feature)

        feature = feature.transpose(0, 1)  # N x 20 x 320
        feature = feature.view(-1, self.linear_input_dim)   # N x (20 x 320)
        feature = torch.stack([
            self.dense(feature[index].unsqueeze(dim=0))
            for index in range(feature.size()[0])])  # N x 80
        feature = feature.transpose(0, 1)  # 80 x N

        return feature

class AttentionConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, groups=1, bias=False):
        super(AttentionConv, self).__init__()
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.groups = groups

        assert self.out_channels % self.groups == 0, "out_channels should be divided by groups. (example: out_channels: 40, groups: 4)"

        self.rel_h = nn.Parameter(torch.randn(out_channels // 2, 1, 1, kernel_size, 1), requires_grad=True)
        self.rel_w = nn.Parameter(torch.randn(out_channels // 2, 1, 1, 1, kernel_size), requires_grad=True)

        self.key_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)
        self.query_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)
        self.value_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)

        self.reset_parameters()

    def forward(self, x):
        batch, channels, height, width = x.size()

        padded_x = F.pad(x, [self.padding, self.padding, self.padding, self.padding])
        q_out = self.query_conv(x)
        k_out = self.key_conv(padded_x)
        v_out = self.value_conv(padded_x)

        k_out = k_out.unfold(2, self.kernel_size, self.stride).unfold(3, self.kernel_size, self.stride)
        v_out = v_out.unfold(2, self.kernel_size, self.stride).unfold(3, self.kernel_size, self.stride)

        k_out_h, k_out_w = k_out.split(self.out_channels // 2, dim=1)
        k_out = torch.cat((k_out_h + self.rel_h, k_out_w + self.rel_w), dim=1)

        k_out = k_out.contiguous().view(batch, self.groups, self.out_channels // self.groups, height, width, -1)
        v_out = v_out.contiguous().view(batch, self.groups, self.out_channels // self.groups, height, width, -1)

        q_out = q_out.view(batch, self.groups, self.out_channels // self.groups, height, width, 1)

        out = q_out * k_out
        out = F.softmax(out, dim=-1)
        out = torch.einsum('bnchwk,bnchwk -> bnchw', out, v_out).view(batch, -1, height, width)

        return out

    def reset_parameters(self):
        init.kaiming_normal_(self.key_conv.weight, mode='fan_out', nonlinearity='relu')
        init.kaiming_normal_(self.value_conv.weight, mode='fan_out', nonlinearity='relu')
        init.kaiming_normal_(self.query_conv.weight, mode='fan_out', nonlinearity='relu')

        init.normal_(self.rel_h, 0, 1)
        init.normal_(self.rel_w, 0, 1)

```

# Experiment Results

2 separate datasets were used for training - the original Librispeech dataset used in the forked repository, and our own scraped TedSrt dataset. This is because the models all performed poorly on our TedSrt dataset, leading us to believe that our dataset was flawed. Training on the Librispeech dataset instead showed a significant difference in performance between models.

### TedSrt

> 200 MB of audio data was scraped from [Ted2Srt](https://ted2srt.org/). We tried training on 1 GB of dataset but the results were roughly similar.

Baseline Model

| Extractor | Classifier | CTC-WER | CTC-Loss |
|-----------|------------|---------|----------|
| CNN       | RNN        | 0.92    | 6.56     |

Comparing performance of Extractor

| Extractor | Classifier | CTC-WER | CTC-Loss |
|-----------|------------|---------|----------|
| MLP       | RNN        | 0.99    | 6.58     |
| ANN       | RNN        | 1       | 6.92     |
| RNN       | RNN        | 0.95    | 6.53     |

Comparing performance of Classifier

| Extractor | Classifier | CTC-WER | CTC-Loss |
|-----------|------------|---------|----------|
| CNN       | MLP        | 0.90    |5.746     |
| CNN       | CNN        | 0.95    |6.745     |
| CNN       | ANN        | 0.99    |6.98      |

---

### LibriSpeech

> This is a public dataset by [OpenSLR](https://www.openslr.org/12/). The repository is already designed to accommodate this dataset. The subset that we used for training is 6 GB in size.

Baseline Model

| Extractor | Classifier | CTC-WER | CTC-Loss |
|-----------|------------|---------|----------|
| CNN       | RNN        | 0.21926 | 1.32976  |

Comparing performance of Extractor

| Extractor | Classifier | CTC-WER | CTC-Loss |
|-----------|------------|---------|----------|
| MLP       | RNN        | 0.99    | 6.63263  |
| ANN       | RNN        | 0.9954  | 6.87     |
| RNN       | RNN        | 0.263   | 1.62     |

Comparing performance of Classifier

| Extractor | Classifier | CTC-WER | CTC-Loss |
|-----------|------------|---------|----------|
| CNN       | MLP        | 0.5884  | 4.092    |
| CNN       | CNN        | 1       | 7.031    |
| CNN       | ANN        | 0.9982  | 7.065    |

# Discussion

### Our TedSrt dataset is likely flawed

All models performed extremely poorly (WER >= 0.90) on the TedSrt dataset. It is possible that insufficient data was scraped to train the models (we tried up to 1 GB of data).

This is a example from the best performing model trained with TedSrt:

Label Text: **AS A KID THE IDEA OF ABORTION MADE ME SO SAD THAT I KNEW IF I**

Prediction Text: **AND I WE IN THE ON**

Clearly the model overfits to the training data as shown from the picture below. The training loss reduce to zero while validation loss remains hight, so as the WER. It only learns the most frequent word and other trivial features from the training set. At the testing time, the model try to predict to the most frequent words.

Alternatively, the quality of the data might not be good enough for training (audio quality, time annotation etc.). Due to this, we use the Librispeech dataset for a models comparison.

---

### All models performed worse than the baseline

For the Librispeech dataset, all models had greater WER and loss than the baseline. This is expected because the BiLSTM used in the baseline has long-short term memory specifically designed for better training speech-related networks, whereas our classifiers do not.

---

### The RNN-RNN and CNN-MLP networks performed quite well

For the RNN extractor, it is possible that there is a significant element of time-correlation in the audio, i.e. the features at timestep $t$ are heavily related to the features at timestep $t-1$, and somewhat related to the features at $t-2$... etc. Therefore, an RNN extractor is still quite capable of summarizing these temporal relationships, albeit less efficiently than the CNN extractor.

MLP classifier is completely local, meaning that it does not look after time frame before and after and only predicts using the information from current timeframe. The fact that it is still able to reduce the WER significantly shows that the correlation between audio and the corresponding word token is strong enough, such that the word can still be identified without temporal information.

On the other hand, for CNN classifier, although the filters can slide over the entire audio sequence, the small filters might not be able to capture the temporal information for such a long time. Meanwhile it also loss the local information, and ends up and perform even worse the MLP.

One other possibility for why the MLP classifier performed better is because more dense layers were used. Only 1 dense layer was used for both the CNN and ANN classifiers, owing to the difficulty of having an unknown number of timesteps. This could have contributed to the success of the MLP over the other NNs.

---

### Attention Neural Networks performed very poorly

It is possible that insufficient Self-Attention layers were used. Due to hardware and compute limitations, only a single layer could fit into RAM. It is possible that, with a hardware upgrade and more Self-Attention layers, the performance could improve.

# Conclusion

This project has found that apart from the CNN-RNN baseline, the RNN-RNN and CNN-MLP models also perform adequately. The proposed reason for the RNN extractor's effectiveness is its ability to capture temporal connections better, and the proposed reason for the MLP is the presence of more dense layers and the strong correlation between the audio features and corresponding word token. However, the CNN-RNN baseline still performed the best, owing to the more sophisticated CNN and RNN networks used.

We conclude that further research on this topic can be conducted by varying the number of layers used in the models, and using that to verify or dispute the results obtained in our experimetns.

# References

<a name="reference"></a>

1. [Speech Processing for Machine Learning: Filter banks, Mel-Frequency Cepstral Coefficients (MFCCs) and What's In-Between](https://haythamfayek.com/2016/04/21/speech-processing-for-machine-learning.html), H M Fayek.

2. [Towards End-to-End Speech Recognition with Deep Convolutional Neural Networks](https://arxiv.org/abs/1701.02720), Y. Zhang *et al*.

3. [Very Deep Convolutional Networks for End-to-End Speech Recognition](https://arxiv.org/abs/1610.03022), Y. Zhang *et al*.

4. [Advances in Joint CTC-Attention based End-to-End Speech Recognition with a Deep CNN Encoder and RNN-LM](https://arxiv.org/abs/1706.02737), T. Hori *et al*.

5. [Stand-Alone Self-Attention in Vision Models](https://arxiv.org/abs/1906.05909), P. Ramachandran *et al*.