{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"%matplotlib inline\nimport numpy as np\nimport pandas as pd\nimport time\nimport os\nimport copy\nimport torch\nimport torch.nn as nn\nimport torchvision\nfrom torchvision import transforms\nimport cv2\nimport librosa\nfrom librosa import display\nfrom pathlib import Path\nimport multiprocessing\nfrom tqdm import tqdm\nfrom matplotlib import gridspec\nfrom datetime import datetime\nimport matplotlib.pyplot as plt\nimport plotly.express as px\nimport seaborn as sns\nsns.set_palette(sns.color_palette('tab20', 20))\nfrom IPython.core.interactiveshell import InteractiveShell\nInteractiveShell.ast_node_interactivity = \"all\"\nfrom sklearn import metrics\nfrom plotly.offline import init_notebook_mode, iplot, plot\ninit_notebook_mode(connected=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"nb_start = datetime.now()\n\nplt.rcParams['figure.figsize'] = [16, 10]\nplt.rcParams['font.size'] = 14\npd.set_option('display.max_columns', 99)\npd.set_option('display.max_rows', 99)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"from dataclasses import dataclass\n\n@dataclass\nclass BirdLists:\n    all = [\n        'aldfly', 'ameavo', 'amebit', 'amecro', 'amegfi', 'amekes', 'amepip', 'amered', 'amerob', 'amewig', 'amewoo',\n        'amtspa', 'annhum', 'astfly', 'baisan', 'baleag', 'balori', 'banswa', 'barswa', 'bawwar', 'belkin1', 'belspa2',\n        'bewwre', 'bkbcuc', 'bkbmag1', 'bkbwar', 'bkcchi', 'bkchum', 'bkhgro', 'bkpwar', 'bktspa', 'blkpho', 'blugrb1',\n        'blujay', 'bnhcow', 'boboli', 'bongul', 'brdowl', 'brebla', 'brespa', 'brncre', 'brnthr', 'brthum', 'brwhaw',\n        'btbwar', 'btnwar', 'btywar', 'buffle', 'buggna', 'buhvir', 'bulori', 'bushti', 'buwtea', 'buwwar', 'cacwre',\n        'calgul', 'calqua', 'camwar', 'cangoo', 'canwar', 'canwre', 'carwre', 'casfin', 'caster1', 'casvir', 'cedwax',\n        'chispa', 'chiswi', 'chswar', 'chukar', 'clanut', 'cliswa', 'comgol', 'comgra', 'comloo', 'commer', 'comnig',\n        'comrav', 'comred', 'comter', 'comyel', 'coohaw', 'coshum', 'cowscj1', 'daejun', 'doccor', 'dowwoo', 'dusfly',\n        'eargre', 'easblu', 'easkin', 'easmea', 'easpho', 'eastow', 'eawpew', 'eucdov', 'eursta', 'evegro', 'fiespa',\n        'fiscro', 'foxspa', 'gadwal', 'gcrfin', 'gnttow', 'gockin', 'gocspa', 'goleag', 'grbher3', 'grcfly', 'greegr',\n        'greroa', 'greyel', 'grhowl', 'grnher', 'grtgra', 'grycat', 'gryfly', 'haiwoo', 'hamfly', 'herthr', 'hoomer',\n        'hoowar', 'horgre', 'horlar', 'houfin', 'houspa', 'houwre', 'indbun', 'juntit1', 'killde', 'labwoo', 'larspa',\n        'lazbun', 'leabit', 'leafly', 'leasan', 'lecthr', 'lesgol', 'lesnig', 'lesyel', 'lewwoo', 'linspa', 'lobcur',\n        'lobdow', 'logshr', 'lotduc', 'louwat', 'macwar', 'magwar', 'mallar3', 'marwre', 'merlin', 'moublu', 'mouchi',\n        'moudov', 'norcar', 'norfli', 'norhar2', 'normoc', 'norpar', 'norpin', 'norsho', 'norwat', 'nrwswa', 'nutwoo',\n        'olsfly', 'orcwar', 'osprey', 'ovenbi1', 'palwar', 'pasfly', 'pecsan', 'perfal', 'phaino', 'pibgre', 'pilwoo',\n        'pingro', 'pinjay', 'pinsis', 'pinwar', 'plsvir', 'prawar', 'purfin', 'pygnut', 'rebmer', 'rebnut', 'rebsap',\n        'rebwoo', 'redcro', 'redhea', 'reevir1', 'renpha', 'reshaw', 'rethaw', 'rewbla', 'ribgul', 'rinduc', 'robgro',\n        'rocpig', 'rocwre', 'rthhum', 'ruckin', 'rudduc', 'rufgro', 'rufhum', 'rusbla', 'sagspa1', 'sagthr', 'savspa',\n        'saypho', 'scatan', 'scoori', 'semplo', 'semsan', 'sheowl', 'shshaw', 'snobun', 'snogoo', 'solsan', 'sonspa',\n        'sora', 'sposan', 'spotow', 'stejay', 'swahaw', 'swaspa', 'swathr', 'treswa', 'truswa', 'tuftit', 'tunswa',\n        'veery', 'vesspa', 'vigswa', 'warvir', 'wesblu', 'wesgre', 'weskin', 'wesmea', 'wessan', 'westan', 'wewpew',\n        'whbnut', 'whcspa', 'whfibi', 'whtspa', 'whtswi', 'wilfly', 'wilsni1', 'wiltur', 'winwre3', 'wlswar', 'wooduc',\n        'wooscj2', 'woothr', 'y00475', 'yebfly', 'yebsap', 'yehbla', 'yelwar', 'yerwar', 'yetvir'\n    ]\n    top_165 = [\n        'amecro', 'amegfi', 'amekes', 'amered', 'amerob', 'amewig', 'amtspa', 'annhum', 'baleag', 'balori', 'barswa',\n        'bawwar', 'belkin1', 'bewwre', 'bkbmag1', 'bkcchi', 'bkhgro', 'blkpho', 'blujay', 'bnhcow', 'bongul', 'brdowl',\n        'brebla', 'brncre', 'brnthr', 'btbwar', 'btnwar', 'buffle', 'buggna', 'buhvir', 'bushti', 'buwtea', 'calgul',\n        'calqua', 'cangoo', 'carwre', 'caster1', 'cedwax', 'chispa', 'chiswi', 'chswar', 'cliswa', 'comgol', 'comgra',\n        'comloo', 'commer', 'comrav', 'comyel', 'coohaw', 'cowscj1', 'daejun', 'doccor', 'dowwoo', 'easblu', 'easkin',\n        'easmea', 'easpho', 'eastow', 'eawpew', 'eucdov', 'eursta', 'fiespa', 'fiscro', 'foxspa', 'gadwal', 'gockin',\n        'gocspa', 'grbher3', 'grcfly', 'greegr', 'greyel', 'grhowl', 'grnher', 'grtgra', 'grycat', 'haiwoo', 'herthr',\n        'hoomer', 'horgre', 'horlar', 'houfin', 'houspa', 'houwre', 'indbun', 'killde', 'leafly', 'leasan', 'lesgol',\n        'lesyel', 'linspa', 'logshr', 'magwar', 'mallar3', 'marwre', 'merlin', 'mouchi', 'moudov', 'norcar', 'norfli',\n        'norhar2', 'normoc', 'norpar', 'norpin', 'norsho', 'nrwswa', 'nutwoo', 'orcwar', 'osprey', 'ovenbi1', 'palwar',\n        'perfal', 'pibgre', 'pilwoo', 'pinsis', 'pinwar', 'purfin', 'rebmer', 'rebnut', 'rebwoo', 'redhea', 'reevir1',\n        'reshaw', 'rethaw', 'rewbla', 'ribgul', 'rinduc', 'robgro', 'rocpig', 'rthhum', 'ruckin', 'rudduc', 'savspa',\n        'saypho', 'scatan', 'semplo', 'shshaw', 'snogoo', 'sonspa', 'sposan', 'spotow', 'stejay', 'swaspa', 'swathr',\n        'treswa', 'tuftit', 'veery', 'vigswa', 'warvir', 'wesblu', 'weskin', 'wesmea', 'westan', 'whbnut', 'whcspa',\n        'whtspa', 'wilsni1', 'wiltur', 'wlswar', 'wooduc', 'woothr', 'y00475', 'yebsap', 'yelwar', 'yerwar', 'yetvir'\n    ]\n    non_zero = [\n        'ameavo', 'amebit', 'amecro', 'amegfi', 'amekes', 'amepip', 'amerob', 'amewig', 'amewoo', 'amtspa', 'annhum',\n        'astfly', 'baisan', 'baleag', 'banswa', 'belkin1', 'belspa2', 'bkbcuc', 'bkbmag1', 'bkbwar', 'bkcchi', 'bkchum',\n        'bkpwar', 'bktspa', 'blkpho', 'blugrb1', 'blujay', 'boboli', 'bongul', 'brdowl', 'brebla', 'brespa', 'brthum',\n        'brwhaw', 'btbwar', 'btnwar', 'btywar', 'buffle', 'buhvir', 'bushti', 'buwtea', 'buwwar', 'cacwre', 'calgul',\n        'calqua', 'camwar', 'cangoo', 'canwar', 'canwre', 'casfin', 'caster1', 'casvir', 'cedwax', 'chiswi', 'chswar',\n        'chukar', 'clanut', 'cliswa', 'comgol', 'commer', 'comnig', 'comred', 'comter', 'comyel', 'coohaw', 'coshum',\n        'cowscj1', 'daejun', 'doccor', 'dusfly', 'eargre', 'easkin', 'easmea', 'eawpew', 'eucdov', 'evegro', 'fiespa',\n        'fiscro', 'gadwal', 'gcrfin', 'gnttow', 'gocspa', 'goleag', 'grcfly', 'greroa', 'greyel', 'grnher', 'grtgra',\n        'grycat', 'gryfly', 'hamfly', 'herthr', 'hoomer', 'hoowar', 'horgre', 'horlar', 'indbun', 'juntit1', 'labwoo',\n        'larspa', 'lazbun', 'leabit', 'leafly', 'leasan', 'lecthr', 'lesgol', 'lesnig', 'lesyel', 'lewwoo', 'linspa',\n        'lobcur', 'lobdow', 'logshr', 'lotduc', 'louwat', 'macwar', 'magwar', 'merlin', 'moublu', 'mouchi', 'norfli',\n        'norhar2', 'norpin', 'norsho', 'norwat', 'nrwswa', 'olsfly', 'ovenbi1', 'palwar', 'pasfly', 'pecsan', 'perfal',\n        'phaino', 'pibgre', 'pingro', 'pinjay', 'pinsis', 'pinwar', 'plsvir', 'prawar', 'purfin', 'pygnut', 'rebmer',\n        'rebsap', 'rebwoo', 'redhea', 'reevir1', 'renpha', 'reshaw', 'rewbla', 'rinduc', 'robgro', 'rocwre', 'rthhum',\n        'rudduc', 'rufgro', 'rufhum', 'rusbla', 'sagspa1', 'sagthr', 'saypho', 'scatan', 'scoori', 'semplo', 'semsan',\n        'sheowl', 'shshaw', 'snobun', 'snogoo', 'solsan', 'sonspa', 'sora', 'sposan', 'stejay', 'swahaw', 'swaspa',\n        'truswa', 'tunswa', 'veery', 'vesspa', 'vigswa', 'wesblu', 'wesgre', 'weskin', 'wessan', 'westan', 'wewpew',\n        'whcspa', 'whfibi', 'whtswi', 'wilfly', 'wilsni1', 'wiltur', 'winwre3', 'wooduc', 'wooscj2', 'woothr', 'y00475',\n        'yebfly', 'yebsap', 'yehbla', 'yetvir'\n    ]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"SAMPLE_RATE = 32000\nCHUNK_DURATION = 5\nBATCH_SIZE = 64\nchunk_len = SAMPLE_RATE * CHUNK_DURATION","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"import numpy as np\nimport librosa\nimport math\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn.parameter import Parameter\n\n\nclass DFTBase(nn.Module):\n    def __init__(self):\n        \"\"\"Base class for DFT and IDFT matrix\"\"\"\n        super(DFTBase, self).__init__()\n\n    def dft_matrix(self, n):\n        (x, y) = np.meshgrid(np.arange(n), np.arange(n))\n        omega = np.exp(-2 * np.pi * 1j / n)\n        W = np.power(omega, x * y)\n        return W\n\n    def idft_matrix(self, n):\n        (x, y) = np.meshgrid(np.arange(n), np.arange(n))\n        omega = np.exp(2 * np.pi * 1j / n)\n        W = np.power(omega, x * y)\n        return W\n\n\nclass DFT(DFTBase):\n    def __init__(self, n, norm):\n        \"\"\"Calculate DFT, IDFT, RDFT, IRDFT. \n        Args:\n          n: fft window size\n          norm: None | 'ortho'\n        \"\"\"\n        super(DFT, self).__init__()\n\n        self.W = self.dft_matrix(n)\n        self.inv_W = self.idft_matrix(n)\n\n        self.W_real = torch.Tensor(np.real(self.W))\n        self.W_imag = torch.Tensor(np.imag(self.W))\n        self.inv_W_real = torch.Tensor(np.real(self.inv_W))\n        self.inv_W_imag = torch.Tensor(np.imag(self.inv_W))\n\n        self.n = n\n        self.norm = norm\n\n    def dft(self, x_real, x_imag):\n        \"\"\"Calculate DFT of signal. \n        Args:\n          x_real: (n,), signal real part\n          x_imag: (n,), signal imag part\n        Returns:\n          z_real: (n,), output real part\n          z_imag: (n,), output imag part\n        \"\"\"\n        z_real = torch.matmul(x_real, self.W_real) - torch.matmul(x_imag, self.W_imag)\n        z_imag = torch.matmul(x_imag, self.W_real) + torch.matmul(x_real, self.W_imag)\n\n        if self.norm is None:\n            pass\n        elif self.norm == 'ortho':\n            z_real /= math.sqrt(self.n)\n            z_imag /= math.sqrt(self.n)\n\n        return z_real, z_imag\n\n        return z_real, z_imag\n\n    def idft(self, x_real, x_imag):\n        \"\"\"Calculate IDFT of signal. \n        Args:\n          x_real: (n,), signal real part\n          x_imag: (n,), signal imag part\n        Returns:\n          z_real: (n,), output real part\n          z_imag: (n,), output imag part\n        \"\"\"\n        z_real = torch.matmul(x_real, self.inv_W_real) - torch.matmul(x_imag, self.inv_W_imag)\n        z_imag = torch.matmul(x_imag, self.inv_W_real) + torch.matmul(x_real, self.inv_W_imag)\n\n        if self.norm is None:\n            z_real /= self.n\n        elif self.norm == 'ortho':\n            z_real /= math.sqrt(n)\n            z_imag /= math.sqrt(n)\n\n        return z_real, z_imag\n\n    def rdft(self, x_real):\n        \"\"\"Calculate right DFT of signal. \n        Args:\n          x_real: (n,), signal real part\n          x_imag: (n,), signal imag part\n        Returns:\n          z_real: (n // 2 + 1,), output real part\n          z_imag: (n // 2 + 1,), output imag part\n        \"\"\"\n        n_rfft = self.n // 2 + 1\n        z_real = torch.matmul(x_real, self.W_real[..., 0 : n_rfft])\n        z_imag = torch.matmul(x_real, self.W_imag[..., 0 : n_rfft])\n\n        if self.norm is None:\n            pass\n        elif self.norm == 'ortho':\n            z_real /= math.sqrt(self.n)\n            z_imag /= math.sqrt(self.n)\n\n        return z_real, z_imag\n\n    def irdft(self, x_real, x_imag):\n        \"\"\"Calculate inverse right DFT of signal. \n        Args:\n          x_real: (n // 2 + 1,), signal real part\n          x_imag: (n // 2 + 1,), signal imag part\n        Returns:\n          z_real: (n,), output real part\n          z_imag: (n,), output imag part\n        \"\"\"\n        n_rfft = self.n // 2 + 1\n\n        flip_x_real = torch.flip(x_real, dims=(-1,))\n        x_real = torch.cat((x_real, flip_x_real[..., 1 : n_rfft - 1]), dim=-1)\n\n        flip_x_imag = torch.flip(x_imag, dims=(-1,))\n        x_imag = torch.cat((x_imag, -1. * flip_x_imag[..., 1 : n_rfft - 1]), dim=-1)\n\n        z_real = torch.matmul(x_real, self.inv_W_real) - torch.matmul(x_imag, self.inv_W_imag)\n\n        if self.norm is None:\n            z_real /= self.n\n        elif self.norm == 'ortho':\n            z_real /= math.sqrt(n)\n\n        return z_real\n        \n\nclass STFT(DFTBase):\n    def __init__(self, n_fft=2048, hop_length=None, win_length=None, \n        window='hann', center=True, pad_mode='reflect', freeze_parameters=True):\n        \"\"\"Implementation of STFT with Conv1d. The function has the same output \n        of librosa.core.stft\n        \"\"\"\n        super(STFT, self).__init__()\n\n        assert pad_mode in ['constant', 'reflect']\n\n        self.n_fft = n_fft\n        self.center = center\n        self.pad_mode = pad_mode\n\n        # By default, use the entire frame\n        if win_length is None:\n            win_length = n_fft\n\n        # Set the default hop, if it's not already specified\n        if hop_length is None:\n            hop_length = int(win_length // 4)\n\n        fft_window = librosa.filters.get_window(window, win_length, fftbins=True)\n\n        # Pad the window out to n_fft size\n        fft_window = librosa.util.pad_center(fft_window, n_fft)\n\n        # DFT & IDFT matrix\n        self.W = self.dft_matrix(n_fft)\n\n        out_channels = n_fft // 2 + 1\n\n        self.conv_real = nn.Conv1d(in_channels=1, out_channels=out_channels, \n            kernel_size=n_fft, stride=hop_length, padding=0, dilation=1, \n            groups=1, bias=False)\n\n        self.conv_imag = nn.Conv1d(in_channels=1, out_channels=out_channels, \n            kernel_size=n_fft, stride=hop_length, padding=0, dilation=1, \n            groups=1, bias=False)\n\n        self.conv_real.weight.data = torch.Tensor(\n            np.real(self.W[:, 0 : out_channels] * fft_window[:, None]).T)[:, None, :]\n        # (n_fft // 2 + 1, 1, n_fft)\n\n        self.conv_imag.weight.data = torch.Tensor(\n            np.imag(self.W[:, 0 : out_channels] * fft_window[:, None]).T)[:, None, :]\n        # (n_fft // 2 + 1, 1, n_fft)\n\n        if freeze_parameters:\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def forward(self, input):\n        \"\"\"input: (batch_size, data_length)\n        Returns:\n          real: (batch_size, n_fft // 2 + 1, time_steps)\n          imag: (batch_size, n_fft // 2 + 1, time_steps)\n        \"\"\"\n\n        x = input[:, None, :]   # (batch_size, channels_num, data_length)\n\n        if self.center:\n            x = F.pad(x, pad=(self.n_fft // 2, self.n_fft // 2), mode=self.pad_mode)\n\n        real = self.conv_real(x)\n        imag = self.conv_imag(x)\n        # (batch_size, n_fft // 2 + 1, time_steps)\n\n        real = real[:, None, :, :].transpose(2, 3)\n        imag = imag[:, None, :, :].transpose(2, 3)\n        # (batch_size, 1, time_steps, n_fft // 2 + 1)\n\n        return real, imag\n\n\ndef magphase(real, imag):\n    mag = (real ** 2 + imag ** 2) ** 0.5\n    cos = real / torch.clamp(mag, 1e-10, np.inf)\n    sin = imag / torch.clamp(mag, 1e-10, np.inf)\n    return mag, cos, sin\n\n\nclass ISTFT(DFTBase):\n    def __init__(self, n_fft=2048, hop_length=None, win_length=None, \n        window='hann', center=True, pad_mode='reflect', freeze_parameters=True):\n        \"\"\"Implementation of ISTFT with Conv1d. The function has the same output \n        of librosa.core.istft\n        \"\"\"\n        super(ISTFT, self).__init__()\n\n        assert pad_mode in ['constant', 'reflect']\n\n        self.n_fft = n_fft\n        self.hop_length = hop_length\n        self.win_length = win_length\n        self.window = window\n        self.center = center\n        self.pad_mode = pad_mode\n\n        # By default, use the entire frame\n        if win_length is None:\n            win_length = n_fft\n\n        # Set the default hop, if it's not already specified\n        if hop_length is None:\n            hop_length = int(win_length // 4)\n\n        # DFT & IDFT matrix\n        self.W = self.idft_matrix(n_fft) / n_fft\n\n        self.conv_real = nn.Conv1d(in_channels=n_fft, out_channels=n_fft, \n            kernel_size=1, stride=1, padding=0, dilation=1, \n            groups=1, bias=False)\n\n        self.conv_imag = nn.Conv1d(in_channels=n_fft, out_channels=n_fft, \n            kernel_size=1, stride=1, padding=0, dilation=1, \n            groups=1, bias=False)\n\n        self.reverse = nn.Conv1d(in_channels=n_fft // 2 + 1, \n            out_channels=n_fft // 2 - 1, kernel_size=1, bias=False)\n\n        self.overlap_add = nn.ConvTranspose2d(in_channels=n_fft, \n            out_channels=1, kernel_size=(n_fft, 1), stride=(self.hop_length, 1), bias=False)\n\n        self.ifft_window_sum = []\n\n        self.init_weights()\n\n        if freeze_parameters:\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def init_weights(self):\n        ifft_window = librosa.filters.get_window(self.window, self.win_length, fftbins=True)\n        \"\"\"(win_length,)\"\"\"\n\n        # Pad the window to n_fft\n        ifft_window = librosa.util.pad_center(ifft_window, self.n_fft)\n\n        self.conv_real.weight.data = torch.Tensor(\n            np.real(self.W * ifft_window[None, :]).T)[:, :, None]\n        # (n_fft // 2 + 1, 1, n_fft)\n\n        self.conv_imag.weight.data = torch.Tensor(\n            np.imag(self.W * ifft_window[None, :]).T)[:, :, None]\n        # (n_fft // 2 + 1, 1, n_fft)\n\n        tmp = np.zeros((self.n_fft // 2 - 1, self.n_fft // 2 + 1, 1))\n        tmp[:, 1 : -1, 0] = np.array(np.eye(self.n_fft // 2 - 1)[::-1])\n        self.reverse.weight.data = torch.Tensor(tmp)\n        \"\"\"(n_fft // 2 - 1, n_fft // 2 + 1, 1)\"\"\"\n\n        self.overlap_add.weight.data = torch.Tensor(np.eye(self.n_fft)[:, None, :, None])\n        \"\"\"(n_fft, 1, n_fft, 1)\"\"\"\n\n    def get_ifft_window(self, n_frames):\n        device = next(self.parameters()).device\n\n        ifft_window_sum = librosa.filters.window_sumsquare(self.window, n_frames,\n            win_length=self.win_length, n_fft=self.n_fft, hop_length=self.hop_length)\n\n        ifft_window_sum = np.clip(ifft_window_sum, 1e-8, np.inf)\n        ifft_window_sum = torch.Tensor(ifft_window_sum).to(device)\n        return ifft_window_sum\n\n    def forward(self, real_stft, imag_stft, length):\n        \"\"\"input: (batch_size, 1, time_steps, n_fft // 2 + 1)\n        Returns:\n          real: (batch_size, data_length)\n        \"\"\"\n        assert real_stft.ndimension() == 4 and imag_stft.ndimension() == 4\n        device = next(self.parameters()).device\n        batch_size = real_stft.shape[0]\n\n        real_stft = real_stft[:, 0, :, :].transpose(1, 2)\n        imag_stft = imag_stft[:, 0, :, :].transpose(1, 2)\n        # (batch_size, n_fft // 2 + 1, time_steps)\n\n        # Full stft, using flip is not supported by ONNX.\n        # full_real_stft = torch.cat((real_stft, torch.flip(real_stft[:, 1 : -1, :], dims=[1])), dim=1)\n        # full_imag_stft = torch.cat((imag_stft, - torch.flip(imag_stft[:, 1 : -1, :], dims=[1])), dim=1)\n        full_real_stft = torch.cat((real_stft, self.reverse(real_stft)), dim=1)\n        full_imag_stft = torch.cat((imag_stft, - self.reverse(imag_stft)), dim=1)\n        \"\"\"(1, n_fft, time_steps)\"\"\"\n\n        # IDFT\n        s_real = self.conv_real(full_real_stft) - self.conv_imag(full_imag_stft)\n        s_real = s_real[..., None]  # (1, n_fft, time_steps, 1)\n        y = self.overlap_add(s_real)[:, 0, :, 0]    # (1, samples_num)\n\n        # Divide window\n        if len(self.ifft_window_sum) != y.shape[1]:\n            frames_num = real_stft.shape[2]\n            self.ifft_window_sum = self.get_ifft_window(frames_num)\n            \n        y = y / self.ifft_window_sum[None, 0 : y.shape[1]]\n\n        # Trim or pad to length\n        if length is None:\n            if self.center:\n                y = y[:, self.n_fft // 2 : -self.n_fft // 2]\n        else:\n            if self.center:\n                start = self.n_fft // 2\n            else:\n                start = 0\n\n            y = y[:, start : start + length]\n            (batch_size, len_y) = y.shape\n            if y.shape[-1] < length:\n                y = torch.cat((y, torch.zeros(batch_size, length - len_y).to(device)), dim=-1)\n        \n        return y\n        \n\nclass Spectrogram(nn.Module):\n    def __init__(self, n_fft=2048, hop_length=None, win_length=None, \n        window='hann', center=True, pad_mode='reflect', power=2.0, \n        freeze_parameters=True):\n        \"\"\"Calculate spectrogram using pytorch. The STFT is implemented with \n        Conv1d. The function has the same output of librosa.core.stft\n        \"\"\"\n        super(Spectrogram, self).__init__()\n\n        self.power = power\n\n        self.stft = STFT(n_fft=n_fft, hop_length=hop_length, \n            win_length=win_length, window=window, center=center, \n            pad_mode=pad_mode, freeze_parameters=True)\n\n    def forward(self, input):\n        \"\"\"input: (batch_size, 1, time_steps, n_fft // 2 + 1)\n        Returns:\n          spectrogram: (batch_size, 1, time_steps, n_fft // 2 + 1)\n        \"\"\"\n\n        (real, imag) = self.stft.forward(input)\n        # (batch_size, n_fft // 2 + 1, time_steps)\n\n        spectrogram = real ** 2 + imag ** 2\n\n        if self.power == 2.0:\n            pass\n        else:\n            spectrogram = spectrogram ** (power / 2.0)\n\n        return spectrogram\n\n\nclass LogmelFilterBank(nn.Module):\n    def __init__(self, sr=32000, n_fft=2048, n_mels=64, fmin=50, fmax=14000, is_log=True, \n        ref=1.0, amin=1e-10, top_db=80.0, freeze_parameters=True):\n        \"\"\"Calculate logmel spectrogram using pytorch. The mel filter bank is \n        the pytorch implementation of as librosa.filters.mel \n        \"\"\"\n        super(LogmelFilterBank, self).__init__()\n\n        self.is_log = is_log\n        self.ref = ref\n        self.amin = amin\n        self.top_db = top_db\n\n        self.melW = librosa.filters.mel(sr=sr, n_fft=n_fft, n_mels=n_mels,\n            fmin=fmin, fmax=fmax).T\n        # (n_fft // 2 + 1, mel_bins)\n\n        self.melW = nn.Parameter(torch.Tensor(self.melW))\n\n        if freeze_parameters:\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def forward(self, input):\n        \"\"\"input: (batch_size, channels, time_steps)\n        \n        Output: (batch_size, time_steps, mel_bins)\n        \"\"\"\n\n        # Mel spectrogram\n        mel_spectrogram = torch.matmul(input, self.melW)\n\n        # Logmel spectrogram\n        if self.is_log:\n            output = self.power_to_db(mel_spectrogram)\n        else:\n            output = mel_spectrogram\n\n        return output\n\n\n    def power_to_db(self, input):\n        \"\"\"Power to db, this function is the pytorch implementation of \n        librosa.core.power_to_lb\n        \"\"\"\n        ref_value = self.ref\n        log_spec = 10.0 * torch.log10(torch.clamp(input, min=self.amin, max=np.inf))\n        log_spec -= 10.0 * np.log10(np.maximum(self.amin, ref_value))\n\n        if self.top_db is not None:\n            if self.top_db < 0:\n                raise ParameterError('top_db must be non-negative')\n            log_spec = torch.clamp(log_spec, min=log_spec.max().item() - self.top_db, max=np.inf)\n\n        return log_spec\n","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass DropStripes(nn.Module):\n    def __init__(self, dim, drop_width, stripes_num):\n        \"\"\"Drop stripes. \n\n        Args:\n          dim: int, dimension along which to drop\n          drop_width: int, maximum width of stripes to drop\n          stripes_num: int, how many stripes to drop\n        \"\"\"\n        super(DropStripes, self).__init__()\n\n        assert dim in [2, 3]    # dim 2: time; dim 3: frequency\n\n        self.dim = dim\n        self.drop_width = drop_width\n        self.stripes_num = stripes_num\n\n    def forward(self, input):\n        \"\"\"input: (batch_size, channels, time_steps, freq_bins)\"\"\"\n\n        assert input.ndimension() == 4\n\n        if self.training is False:\n            return input\n\n        else:\n            batch_size = input.shape[0]\n            total_width = input.shape[self.dim]\n\n            for n in range(batch_size):\n                self.transform_slice(input[n], total_width)\n\n            return input\n\n\n    def transform_slice(self, e, total_width):\n        \"\"\"e: (channels, time_steps, freq_bins)\"\"\"\n\n        for _ in range(self.stripes_num):\n            distance = torch.randint(low=0, high=self.drop_width, size=(1,))[0]\n            bgn = torch.randint(low=0, high=total_width - distance, size=(1,))[0]\n\n            if self.dim == 2:\n                e[:, bgn : bgn + distance, :] = 0\n            elif self.dim == 3:\n                e[:, :, bgn : bgn + distance] = 0\n\n\nclass SpecAugmentation(nn.Module):\n    def __init__(self, time_drop_width, time_stripes_num, freq_drop_width, \n        freq_stripes_num):\n        \"\"\"Spec augmetation. \n        [ref] Park, D.S., Chan, W., Zhang, Y., Chiu, C.C., Zoph, B., Cubuk, E.D. \n        and Le, Q.V., 2019. Specaugment: A simple data augmentation method \n        for automatic speech recognition. arXiv preprint arXiv:1904.08779.\n\n        Args:\n          time_drop_width: int\n          time_stripes_num: int\n          freq_drop_width: int\n          freq_stripes_num: int\n        \"\"\"\n\n        super(SpecAugmentation, self).__init__()\n\n        self.time_dropper = DropStripes(dim=2, drop_width=time_drop_width, \n            stripes_num=time_stripes_num)\n\n        self.freq_dropper = DropStripes(dim=3, drop_width=freq_drop_width, \n            stripes_num=freq_stripes_num)\n\n    def forward(self, input):\n        x = self.time_dropper(input)\n        x = self.freq_dropper(x)\n        return x\n","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n# from torchlibrosa.stft import Spectrogram, LogmelFilterBank\n# from torchlibrosa.augmentation import SpecAugmentation\n\n\ndef do_mixup(x, mixup_lambda):\n    \"\"\"Mixup x of even indexes (0, 2, 4, ...) with x of odd indexes\n    (1, 3, 5, ...).\n\n    Args:\n      x: (batch_size * 2, ...)\n      mixup_lambda: (batch_size * 2,)\n\n    Returns:\n      out: (batch_size, ...)\n    \"\"\"\n    out = (x[0:: 2].transpose(0, -1) * mixup_lambda[0:: 2] + \\\n           x[1:: 2].transpose(0, -1) * mixup_lambda[1:: 2]).transpose(0, -1)\n    return out\n\n\ndef interpolate(x, ratio):\n    \"\"\"Interpolate data in time domain. This is used to compensate the\n    resolution reduction in downsampling of a CNN.\n\n    Args:\n      x: (batch_size, time_steps, classes_num)\n      ratio: int, ratio to interpolate\n\n    Returns:\n      upsampled: (batch_size, time_steps * ratio, classes_num)\n    \"\"\"\n    (batch_size, time_steps, classes_num) = x.shape\n    upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1)\n    upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num)\n    return upsampled\n\n\ndef pad_framewise_output(framewise_output, frames_num):\n    \"\"\"Pad framewise_output to the same length as input frames. The pad value\n    is the same as the value of the last frame.\n\n    Args:\n      framewise_output: (batch_size, frames_num, classes_num)\n      frames_num: int, number of frames to pad\n\n    Outputs:\n      output: (batch_size, frames_num, classes_num)\n    \"\"\"\n    pad = framewise_output[:, -1:, :].repeat(1, frames_num - framewise_output.shape[1], 1)\n    \"\"\"tensor for padding\"\"\"\n\n    output = torch.cat((framewise_output, pad), dim=1)\n    \"\"\"(batch_size, frames_num, classes_num)\"\"\"\n\n    return output\n\n\ndef init_layer(layer):\n    \"\"\"Initialize a Linear or Convolutional layer. \"\"\"\n    nn.init.xavier_uniform_(layer.weight)\n\n    if hasattr(layer, 'bias'):\n        if layer.bias is not None:\n            layer.bias.data.fill_(0.)\n\n\ndef init_bn(bn):\n    \"\"\"Initialize a Batchnorm layer. \"\"\"\n    bn.bias.data.fill_(0.)\n    bn.weight.data.fill_(1.)\n\n\nclass ConvBlock(nn.Module):\n    def __init__(self, in_channels, out_channels):\n\n        super(ConvBlock, self).__init__()\n\n        self.conv1 = nn.Conv2d(in_channels=in_channels,\n                               out_channels=out_channels,\n                               kernel_size=(3, 3), stride=(1, 1),\n                               padding=(1, 1), bias=False)\n\n        self.conv2 = nn.Conv2d(in_channels=out_channels,\n                               out_channels=out_channels,\n                               kernel_size=(3, 3), stride=(1, 1),\n                               padding=(1, 1), bias=False)\n\n        self.bn1 = nn.BatchNorm2d(out_channels)\n        self.bn2 = nn.BatchNorm2d(out_channels)\n\n        self.init_weight()\n\n    def init_weight(self):\n        init_layer(self.conv1)\n        init_layer(self.conv2)\n        init_bn(self.bn1)\n        init_bn(self.bn2)\n\n    def forward(self, input, pool_size=(2, 2), pool_type='avg'):\n\n        x = input\n        x = F.relu_(self.bn1(self.conv1(x)))\n        x = F.relu_(self.bn2(self.conv2(x)))\n        if pool_type == 'max':\n            x = F.max_pool2d(x, kernel_size=pool_size)\n        elif pool_type == 'avg':\n            x = F.avg_pool2d(x, kernel_size=pool_size)\n        elif pool_type == 'avg+max':\n            x1 = F.avg_pool2d(x, kernel_size=pool_size)\n            x2 = F.max_pool2d(x, kernel_size=pool_size)\n            x = x1 + x2\n        else:\n            raise Exception('Incorrect argument!')\n\n        return x\n\n\nclass ConvBlock5x5(nn.Module):\n    def __init__(self, in_channels, out_channels):\n\n        super(ConvBlock5x5, self).__init__()\n\n        self.conv1 = nn.Conv2d(in_channels=in_channels,\n                               out_channels=out_channels,\n                               kernel_size=(5, 5), stride=(1, 1),\n                               padding=(2, 2), bias=False)\n\n        self.bn1 = nn.BatchNorm2d(out_channels)\n\n        self.init_weight()\n\n    def init_weight(self):\n        init_layer(self.conv1)\n        init_bn(self.bn1)\n\n    def forward(self, input, pool_size=(2, 2), pool_type='avg'):\n\n        x = input\n        x = F.relu_(self.bn1(self.conv1(x)))\n        if pool_type == 'max':\n            x = F.max_pool2d(x, kernel_size=pool_size)\n        elif pool_type == 'avg':\n            x = F.avg_pool2d(x, kernel_size=pool_size)\n        elif pool_type == 'avg+max':\n            x1 = F.avg_pool2d(x, kernel_size=pool_size)\n            x2 = F.max_pool2d(x, kernel_size=pool_size)\n            x = x1 + x2\n        else:\n            raise Exception('Incorrect argument!')\n\n        return x\n\n\nclass AttBlock(nn.Module):\n    def __init__(self, n_in, n_out, activation='linear', temperature=1.):\n        super(AttBlock, self).__init__()\n\n        self.activation = activation\n        self.temperature = temperature\n        self.att = nn.Conv1d(in_channels=n_in, out_channels=n_out, kernel_size=1, stride=1, padding=0, bias=True)\n        self.cla = nn.Conv1d(in_channels=n_in, out_channels=n_out, kernel_size=1, stride=1, padding=0, bias=True)\n\n        self.bn_att = nn.BatchNorm1d(n_out)\n        self.init_weights()\n\n    def init_weights(self):\n        init_layer(self.att)\n        init_layer(self.cla)\n        init_bn(self.bn_att)\n\n    def forward(self, x):\n        # x: (n_samples, n_in, n_time)\n        norm_att = torch.softmax(torch.clamp(self.att(x), -10, 10), dim=-1)\n        cla = self.nonlinear_transform(self.cla(x))\n        x = torch.sum(norm_att * cla, dim=2)\n        return x, norm_att, cla\n\n    def nonlinear_transform(self, x):\n        if self.activation == 'linear':\n            return x\n        elif self.activation == 'sigmoid':\n            return torch.sigmoid(x)\n\n\nclass Cnn14(nn.Module):\n    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,\n                 fmax, classes_num):\n\n        super(Cnn14, self).__init__()\n\n        window = 'hann'\n        center = True\n        pad_mode = 'reflect'\n        ref = 1.0\n        amin = 1e-10\n        top_db = None\n        \n        self.dataset_mean = 0.\n        self.dataset_std = 1.\n\n        # Spectrogram extractor\n        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size,\n                                                 win_length=window_size, window=window, center=center,\n                                                 pad_mode=pad_mode,\n                                                 freeze_parameters=True)\n\n        # Logmel feature extractor\n        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size,\n                                                 n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin,\n                                                 top_db=top_db,\n                                                 freeze_parameters=True)\n\n        # Spec augmenter\n        self.spec_augmenter = SpecAugmentation(time_drop_width=32, time_stripes_num=2,\n                                               freq_drop_width=8, freq_stripes_num=2)\n\n        self.bn0 = nn.BatchNorm2d(64)\n\n        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)\n        self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)\n        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)\n        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)\n        self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024)\n        self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048)\n\n        self.fc1 = nn.Linear(2048, 2048, bias=True)\n        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)\n\n        self.init_weight()\n\n    def init_weight(self):\n        init_bn(self.bn0)\n        init_layer(self.fc1)\n        init_layer(self.fc_audioset)\n\n    def forward(self, input, mixup_lambda=None):\n        \"\"\"\n        Input: (batch_size, data_length)\"\"\"\n\n        x = self.spectrogram_extractor(input)  # (batch_size, 1, time_steps, freq_bins)\n        x = self.logmel_extractor(x)  # (batch_size, 1, time_steps, mel_bins)\n        \n        x = (x - self.dataset_mean) / self.dataset_std\n        \n        x = x.transpose(1, 3)\n        x = self.bn0(x)\n        x = x.transpose(1, 3)\n\n        if self.training:\n            x = self.spec_augmenter(x)\n\n        # Mixup on spectrogram\n        if self.training and mixup_lambda is not None:\n            x = do_mixup(x, mixup_lambda)\n\n        x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')\n        x = F.dropout(x, p=0.2, training=self.training)\n        x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')\n        x = F.dropout(x, p=0.2, training=self.training)\n        x = self.conv_block3(x, pool_size=(2, 2), pool_type='avg')\n        x = F.dropout(x, p=0.2, training=self.training)\n        x = self.conv_block4(x, pool_size=(2, 2), pool_type='avg')\n        x = F.dropout(x, p=0.2, training=self.training)\n        x = self.conv_block5(x, pool_size=(2, 2), pool_type='avg')\n        x = F.dropout(x, p=0.2, training=self.training)\n        x = self.conv_block6(x, pool_size=(1, 1), pool_type='avg')\n        x = F.dropout(x, p=0.2, training=self.training)\n        x = torch.mean(x, dim=3)\n\n        (x1, _) = torch.max(x, dim=2)\n        x2 = torch.mean(x, dim=2)\n        x = x1 + x2\n        x = F.dropout(x, p=0.5, training=self.training)\n        x = F.relu_(self.fc1(x))\n        embedding = F.dropout(x, p=0.5, training=self.training)\n        clipwise_output = torch.sigmoid(self.fc_audioset(x))\n\n        output_dict = {'clipwise_output': clipwise_output, 'embedding': embedding}\n\n        return output_dict","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def _resnet_conv3x3(in_planes, out_planes):\n    # 3x3 convolution with padding\n    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=1,\n                     padding=1, groups=1, bias=False, dilation=1)\n\ndef _resnet_conv1x1(in_planes, out_planes):\n    # 1x1 convolution\n    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, bias=False)\n\nclass _ResnetBasicBlock(nn.Module):\n    expansion = 1\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,\n                 base_width=64, dilation=1, norm_layer=None):\n        super(_ResnetBasicBlock, self).__init__()\n        if norm_layer is None:\n            norm_layer = nn.BatchNorm2d\n        if groups != 1 or base_width != 64:\n            raise ValueError('_ResnetBasicBlock only supports groups=1 and base_width=64')\n        if dilation > 1:\n            raise NotImplementedError(\"Dilation > 1 not supported in _ResnetBasicBlock\")\n        # Both self.conv1 and self.downsample layers downsample the input when stride != 1\n\n        self.stride = stride\n\n        self.conv1 = _resnet_conv3x3(inplanes, planes)\n        self.bn1 = norm_layer(planes)\n        self.relu = nn.ReLU(inplace=True)\n        self.conv2 = _resnet_conv3x3(planes, planes)\n        self.bn2 = norm_layer(planes)\n        self.downsample = downsample\n        self.stride = stride\n\n        self.init_weights()\n\n    def init_weights(self):\n        init_layer(self.conv1)\n        init_bn(self.bn1)\n        init_layer(self.conv2)\n        init_bn(self.bn2)\n        nn.init.constant_(self.bn2.weight, 0)\n\n    def forward(self, x):\n        identity = x\n\n        if self.stride == 2:\n            out = F.avg_pool2d(x, kernel_size=(2, 2))\n        else:\n            out = x\n\n        out = self.conv1(out)\n        out = self.bn1(out)\n        out = self.relu(out)\n        out = F.dropout(out, p=0.1, training=self.training)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n\n        if self.downsample is not None:\n            identity = self.downsample(identity)\n\n        out += identity\n        out = self.relu(out)\n\n        return out\n\n\nclass _ResnetBottleneck(nn.Module):\n    expansion = 4\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,\n                 base_width=64, dilation=1, norm_layer=None):\n        super(_ResnetBottleneck, self).__init__()\n        if norm_layer is None:\n            norm_layer = nn.BatchNorm2d\n        width = int(planes * (base_width / 64.)) * groups\n        self.stride = stride\n        # Both self.conv2 and self.downsample layers downsample the input when stride != 1\n        self.conv1 = _resnet_conv1x1(inplanes, width)\n        self.bn1 = norm_layer(width)\n        self.conv2 = _resnet_conv3x3(width, width)\n        self.bn2 = norm_layer(width)\n        self.conv3 = _resnet_conv1x1(width, planes * self.expansion)\n        self.bn3 = norm_layer(planes * self.expansion)\n        self.relu = nn.ReLU(inplace=True)\n        self.downsample = downsample\n        self.stride = stride\n\n        self.init_weights()\n\n    def init_weights(self):\n        init_layer(self.conv1)\n        init_bn(self.bn1)\n        init_layer(self.conv2)\n        init_bn(self.bn2)\n        init_layer(self.conv3)\n        init_bn(self.bn3)\n        nn.init.constant_(self.bn3.weight, 0)\n\n    def forward(self, x):\n        identity = x\n\n        if self.stride == 2:\n            x = F.avg_pool2d(x, kernel_size=(2, 2))\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n        out = self.relu(out)\n        out = F.dropout(out, p=0.1, training=self.training)\n\n        out = self.conv3(out)\n        out = self.bn3(out)\n\n        if self.downsample is not None:\n            identity = self.downsample(identity)\n\n        out += identity\n        out = self.relu(out)\n\n        return out\n\nclass _ResNet(nn.Module):\n    def __init__(self, block, layers, zero_init_residual=False,\n                 groups=1, width_per_group=64, replace_stride_with_dilation=None,\n                 norm_layer=None):\n        super(_ResNet, self).__init__()\n\n        if norm_layer is None:\n            norm_layer = nn.BatchNorm2d\n        self._norm_layer = norm_layer\n\n        self.inplanes = 64\n        self.dilation = 1\n        if replace_stride_with_dilation is None:\n            # each element in the tuple indicates if we should replace\n            # the 2x2 stride with a dilated convolution instead\n            replace_stride_with_dilation = [False, False, False]\n        if len(replace_stride_with_dilation) != 3:\n            raise ValueError(\"replace_stride_with_dilation should be None \"\n                             \"or a 3-element tuple, got {}\".format(replace_stride_with_dilation))\n        self.groups = groups\n        self.base_width = width_per_group\n\n        self.layer1 = self._make_layer(block, 64, layers[0], stride=1)\n        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,\n                                       dilate=replace_stride_with_dilation[0])\n        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,\n                                       dilate=replace_stride_with_dilation[1])\n        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,\n                                       dilate=replace_stride_with_dilation[2])\n\n    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):\n        norm_layer = self._norm_layer\n        downsample = None\n        previous_dilation = self.dilation\n        if dilate:\n            self.dilation *= stride\n            stride = 1\n        if stride != 1 or self.inplanes != planes * block.expansion:\n            if stride == 1:\n                downsample = nn.Sequential(\n                    _resnet_conv1x1(self.inplanes, planes * block.expansion),\n                    norm_layer(planes * block.expansion),\n                )\n                init_layer(downsample[0])\n                init_bn(downsample[1])\n            elif stride == 2:\n                downsample = nn.Sequential(\n                    nn.AvgPool2d(kernel_size=2),\n                    _resnet_conv1x1(self.inplanes, planes * block.expansion),\n                    norm_layer(planes * block.expansion),\n                )\n                init_layer(downsample[1])\n                init_bn(downsample[2])\n\n        layers = []\n        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,\n                            self.base_width, previous_dilation, norm_layer))\n        self.inplanes = planes * block.expansion\n        for _ in range(1, blocks):\n            layers.append(block(self.inplanes, planes, groups=self.groups,\n                                base_width=self.base_width, dilation=self.dilation,\n                                norm_layer=norm_layer))\n\n        return nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n\n        return x\n\nclass ResNet38(nn.Module):\n    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,\n                 fmax, classes_num):\n\n        super(ResNet38, self).__init__()\n\n        window = 'hann'\n        center = True\n        pad_mode = 'reflect'\n        ref = 1.0\n        amin = 1e-10\n        top_db = None\n\n        # Spectrogram extractor\n        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size,\n                                                 win_length=window_size, window=window, center=center,\n                                                 pad_mode=pad_mode,\n                                                 freeze_parameters=True)\n\n        # Logmel feature extractor\n        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size,\n                                                 n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin,\n                                                 top_db=top_db,\n                                                 freeze_parameters=True)\n\n        # Spec augmenter\n        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2,\n                                               freq_drop_width=8, freq_stripes_num=2)\n\n        self.bn0 = nn.BatchNorm2d(64)\n\n        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)\n        # self.conv_block2 = ConvBlock(in_channels=64, out_channels=64)\n\n        self.resnet = _ResNet(block=_ResnetBasicBlock, layers=[3, 4, 6, 3], zero_init_residual=True)\n\n        self.conv_block_after1 = ConvBlock(in_channels=512, out_channels=2048)\n\n        self.fc1 = nn.Linear(2048, 2048)\n        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)\n\n        self.init_weights()\n\n    def init_weights(self):\n        init_bn(self.bn0)\n        init_layer(self.fc1)\n        init_layer(self.fc_audioset)\n\n    def forward(self, input, mixup_lambda=None):\n        \"\"\"\n        Input: (batch_size, data_length)\"\"\"\n\n        x = self.spectrogram_extractor(input)  # (batch_size, 1, time_steps, freq_bins)\n        x = self.logmel_extractor(x)  # (batch_size, 1, time_steps, mel_bins)\n\n        x = x.transpose(1, 3)\n        x = self.bn0(x)\n        x = x.transpose(1, 3)\n\n        if self.training:\n            x = self.spec_augmenter(x)\n\n        # Mixup on spectrogram\n        if self.training and mixup_lambda is not None:\n            x = do_mixup(x, mixup_lambda)\n\n        x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')\n        x = F.dropout(x, p=0.2, training=self.training, inplace=True)\n        x = self.resnet(x)\n        x = F.avg_pool2d(x, kernel_size=(2, 2))\n        x = F.dropout(x, p=0.2, training=self.training, inplace=True)\n        x = self.conv_block_after1(x, pool_size=(1, 1), pool_type='avg')\n        x = F.dropout(x, p=0.2, training=self.training, inplace=True)\n        x = torch.mean(x, dim=3)\n\n        (x1, _) = torch.max(x, dim=2)\n        x2 = torch.mean(x, dim=2)\n        x = x1 + x2\n        x = F.dropout(x, p=0.5, training=self.training)\n        x = F.relu_(self.fc1(x))\n        embedding = F.dropout(x, p=0.5, training=self.training)\n        clipwise_output = torch.sigmoid(self.fc_audioset(x))\n\n        output_dict = {'clipwise_output': clipwise_output, 'embedding': embedding}\n\n        return output_dict","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"_window = 'hann'\n_center = True\n_pad_mode = 'reflect'\n_ref = 1.0\n_amin = 1e-10\n_top_db = None\n_model_config = {\n    \"sample_rate\": 32000,\n    \"window_size\": 1024,  # original\n    \"hop_size\": 320,  # original\n    \"mel_bins\": 64,  # original\n    \"fmin\": 50,\n    \"fmax\": 14000,\n    \"classes_num\": 527  # original\n}\n\ndef my_cnn14(n_fft, n_mels, n_classes=100, hop_size=320, fmin=160, fmax=10300):\n    model = Cnn14(**_model_config)\n    model.fc_audioset = nn.Linear(2048, n_classes, bias=True)\n    init_layer(model.fc_audioset)\n    model.spectrogram_extractor = Spectrogram(\n        n_fft=n_fft, hop_length=hop_size, win_length=n_fft, window=_window,\n        center=_center, pad_mode=_pad_mode, freeze_parameters=True)\n    model.logmel_extractor = LogmelFilterBank(\n        sr=SAMPLE_RATE, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax,\n        ref=_ref, amin=_amin, top_db=_top_db, freeze_parameters=True)\n    model.bn0 = nn.BatchNorm2d(n_mels)\n    init_bn(model.bn0)\n    return model\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def read_audio_fast(path):\n    clip, sr_native = librosa.core.audio.__audioread_load(\n        path, offset=0.0, duration=None, dtype=np.float32)\n    clip = librosa.to_mono(clip)\n    if sr_native > 0:\n        clip = librosa.resample(clip, sr_native, SAMPLE_RATE, res_type='kaiser_fast')\n    return clip, SAMPLE_RATE\n\ndef random_chunk(y, duration, sr=SAMPLE_RATE):\n    sample_size = int(duration * sr)\n    if y.shape[0] < sample_size:\n        raise ValueError('Too short clip')\n    t = np.random.randint(0, y.shape[0] - sample_size)\n    return y[t: t + sample_size]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if os.path.exists('/kaggle/input/birdsong-recognition/test_audio'):\n    BASE_TEST_DIR = '/kaggle/input/birdsong-recognition'\n    TEST_AUDIO_DIR = '/kaggle/input/birdsong-recognition/test_audio'\nelse:\n    BASE_TEST_DIR = '/kaggle/input/cornellfaketest'\n    TEST_AUDIO_DIR = '/kaggle/input/cornellfaketest'\nprint(BASE_TEST_DIR, TEST_AUDIO_DIR)\ndf_test = pd.read_csv(f'{BASE_TEST_DIR}/test.csv')\ndf_test.shape\ndf_test.head()\ndf_test.groupby('site').nunique()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"audio_site = {a: s for a, s in df_test[['audio_id', 'site']].drop_duplicates().values}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_model_predictions_for_clip_fast(y, model, device, classes):\n    duration = y.shape[0] // SAMPLE_RATE\n    batch = []\n    start_seconds = []\n    for start in range(0, duration - CHUNK_DURATION + 1, 5):\n        end = start + 5\n        start_seconds.append(start)\n        chunk = y[start * SAMPLE_RATE: end * SAMPLE_RATE]\n        if len(chunk) != CHUNK_DURATION * SAMPLE_RATE:\n            print(chunk.shape)\n            break\n        batch.append(chunk)\n    batch = np.asarray(batch)\n    tensors = torch.from_numpy(batch)\n    tensors = tensors.to(device)\n    with torch.no_grad():\n        preds = model(tensors)['clipwise_output']\n    test_preds = preds.cpu().numpy()\n    pred_df = pd.DataFrame(test_preds, columns=classes)\n    pred_df['start_second'] = start_seconds\n    return pred_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sites = ['site_1', 'site_2', 'site_3']\naudio_ids_by_site = {}\nfor s in sites:\n    audio_ids_by_site[s] = list(df_test[df_test.site == s].audio_id.unique())\n    print(s, len(audio_ids_by_site[s]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Site 1 & 2"},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls /kaggle/input/cornellmodels/","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"models = [\n#     (\n#         'models_CNN14_CV0_C262_M2_5139_13143345.pth',\n#         {'n_fft': 1536, 'hop_size': 320, 'n_mels': 128, 'n_classes': 262, 'fmin': 160, 'fmax': 12000}\n#     ),\n#     (\n#         'models_CNN14_CV0_C262_F18147_05073327.pth',\n#         {'n_fft': 1536, 'hop_size': 360, 'n_mels': 128, 'n_classes': 262, 'fmin': 500, 'fmax': 11000}\n#     ),\n#     (\n#         'models_CNN14_CV0_C165_M2_N0_5816_13174140.pth', \n#         {'n_fft': 1536, 'hop_size': 320, 'n_mels': 128, 'n_classes': 165, 'fmin': 300, 'fmax': 11000}\n#     ),\n    (\n        'models_CNN14_CV1_C262_M1_N1_7045_14005959.pth', \n        {'n_fft': 1536, 'hop_size': 320, 'n_mels': 128, 'n_classes': 262, 'fmin': 300, 'fmax': 11000}\n    ),\n    (\n        'models_CNN14_CV3_C165_M3_N1_F18256_14102207.pth', \n        {'n_fft': 1536, 'hop_size': 320, 'n_mels': 128, 'n_classes': 165, 'fmin': 300, 'fmax': 11000}\n    ),\n    (\n        'models_CNN14_CV3_C262_M0_N1_7063_14141546.pth', \n        {'n_fft': 1536, 'hop_size': 320, 'n_mels': 128, 'n_classes': 262, 'fmin': 300, 'fmax': 11000}\n    )\n]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n                      \ndef load_model(model_name, model_config):\n    model = my_cnn14(**model_config)\n    model.load_state_dict(torch.load(f'/kaggle/input/cornellmodels/{model_name}', map_location=device))\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_spec_stats(ys, model):\n    test_means = []\n    test_stds = []\n    for y in tqdm(ys):\n        duration = y.shape[0] // SAMPLE_RATE\n        batch = []\n        for start in range(0, duration, 5):\n            end = start + 5\n            chunk = y[start * SAMPLE_RATE: end * SAMPLE_RATE]\n            batch.append(chunk)\n        batch = np.array(batch)\n        batch = torch.from_numpy(batch)\n        tensors = batch.to(device)\n        specs = model.spectrogram_extractor(tensors)\n        specs = model.logmel_extractor(specs)\n        specs = specs.cpu().numpy()\n        for spec in specs:\n            test_means.append(spec.mean())\n            test_stds.append(spec.std())\n    test_mean = np.mean(test_means)\n    test_std = np.mean(test_stds)\n    print(f'Test stats: {test_mean:.3} +- {test_std:.3}')\n    return test_mean, test_std","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"site12_preds = []\n\nfor site in ['site_1', 'site_2']:\n    test_audios = {}\n    for audio_id in tqdm(audio_ids_by_site[site]):\n        path = os.path.join(TEST_AUDIO_DIR, f'{audio_id}.mp3')\n        y, _ = read_audio_fast(path)\n        test_audios[audio_id] = y\n    \n    for model_name, model_config in models:\n        model = load_model(model_name, model_config)\n        model = model.to(device)\n        _ = model.eval()\n\n        if '_C262_' in model_name:\n            classes = BirdLists.all\n        elif '_C165_' in model_name:\n            classes = BirdLists.top_165\n        elif '_C202_' in model_name:\n            classes = BirdLists.non_zero\n        \n        if '_N1_' in model_name:\n            test_mean, test_std = get_spec_stats(test_audios.values(), model)\n            model.dataset_mean = test_mean\n            model.dataset_std = test_std\n        \n        for audio_id, y in tqdm(test_audios.items()):\n            pred_df = get_model_predictions_for_clip_fast(y, model, device, classes)\n            pred_df['audio_id'] = audio_id\n            pred_df['site'] = site\n            site12_preds.append(pred_df)\npredictions_12 = pd.concat(site12_preds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"predictions_12.shape\npredictions_12.head()\npredictions_12.tail()\npredictions_12.count()\n","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"zero = [\n    'redcro', 'norcar', 'houwre', 'spotow', 'bewwre', 'carwre', 'swathr', 'comrav', 'warvir', 'whbnut', 'marwre',\n    'normoc', 'wesmea', 'amered', 'whtspa', 'houfin', 'ruckin', 'eastow', 'tuftit', 'chispa', 'grhowl', 'savspa',\n    'wlswar', 'buggna', 'brncre', 'aldfly', 'orcwar', 'bkhgro', 'bnhcow', 'houspa', 'yerwar', 'haiwoo', 'dowwoo',\n    'bawwar', 'foxspa', 'balori', 'rebnut', 'eursta', 'gockin', 'killde', 'comgra', 'barswa', 'osprey', 'mallar3',\n    'pilwoo', 'moudov', 'easpho', 'comloo', 'treswa', 'easblu', 'rethaw', 'grbher3', 'yelwar', 'greegr', 'ribgul',\n    'nutwoo', 'rocpig', 'brnthr', 'norpar', 'bulori',\n]\n\ncommon = ['whcspa', 'grycat', 'herthr', 'amerob', 'reevir1', 'comyel', 'norfli', 'amecro']\nrare = [\n    'doccor', 'baleag', 'buffle', 'amekes', 'annhum', 'reshaw', 'coohaw', 'lesgol', 'norsho', 'commer', 'norhar2',\n    'rinduc', 'amewig', 'blkpho', 'hoomer', 'comgol', 'cowscj1', 'wiltur', 'buwtea', 'rudduc', 'bushti', 'rebmer',\n    'stejay', 'norpin', 'vigswa', 'gocspa', 'redhea', 'wesblu', 'shshaw', 'horgre', 'calqua', 'perfal', 'labwoo',\n    'wewpew', 'bkchum', 'truswa', 'lotduc', 'brwhaw', 'pasfly', 'swahaw', 'eargre', 'whfibi', 'tunswa', 'phaino',\n    'brthum', 'rufgro', 'wooscj2', 'cacwre', 'rebsap', 'whtswi', 'goleag', 'snobun', 'renpha', 'coshum', 'baisan',\n    'sheowl', 'lewwoo', 'lesnig', 'gcrfin', 'sagspa1', 'chukar', 'lecthr',\n]\n\nBIRD_THRESHOLD = 0.33\nTHRESHOLDS = {c: BIRD_THRESHOLD for c in BirdLists.all}\nfor c in zero:\n    THRESHOLDS[c] = 0.6\nfor c in common:\n    THRESHOLDS[c] = 0.3\nfor c in rare:\n    THRESHOLDS[c] = 0.55\n\nTHRESHOLDS","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mean_predictions_12 = predictions_12.fillna(0).groupby(['site', 'audio_id', 'start_second']).mean().reset_index()\nmean_predictions_12['birdmax'] = mean_predictions_12[BirdLists.all].max(axis=1)\nmean_predictions_12['seconds'] = mean_predictions_12.start_second + CHUNK_DURATION\nmean_predictions_12['row_id'] = mean_predictions_12.site + '_' + mean_predictions_12.audio_id + '_' + mean_predictions_12.seconds.astype(str)\nmulti_birds = []\nfor _, row in mean_predictions_12[BirdLists.all].iterrows():\n    guesses = ' '.join([bird for bird, p in row.items() if p >= THRESHOLDS[bird]])\n    multi_birds.append(guesses)\nmean_predictions_12['bird_guesses'] = multi_birds\nmean_predictions_12.shape\nmean_predictions_12.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mean_predictions_12[BirdLists.all].max(axis=1).hist()\nnp.mean(mean_predictions_12.birdmax > 0.5)\nnp.mean(mean_predictions_12.birdmax > 0.4)\nnp.mean(mean_predictions_12.birdmax > 0.33)\nnp.mean(mean_predictions_12.birdmax > 0.3)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"site12_birds = mean_predictions_12.loc[mean_predictions_12.bird_guesses.str.len() > 1, ['row_id', 'bird_guesses']]\nsite12_birds.shape\nsite12_birds.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Site 3"},{"metadata":{"trusted":true},"cell_type":"code","source":"site3_preds = []\nsite = 'site_3'\ntest_audios = {}\nfor audio_id in tqdm(audio_ids_by_site[site]):\n    path = os.path.join(TEST_AUDIO_DIR, f'{audio_id}.mp3')\n    y, _ = read_audio_fast(path)\n    test_audios[audio_id] = y\n\nfor model_name, model_config in models[:1]:\n    model = load_model(model_name, model_config)\n    model = model.to(device)\n    _ = model.eval()\n    classes = BirdLists.all\n    for audio_id, y in tqdm(test_audios.items()):\n        pred_df = get_model_predictions_for_clip_fast(y, model, device, classes)\n        pred_df['audio_id'] = audio_id\n        pred_df['site'] = site\n        site3_preds.append(pred_df)\npredictions_3 = pd.concat(site3_preds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\npredictions_3['birdmax'] = predictions_3[BirdLists.all].max(axis=1)\npredictions_3['bird_guess'] = [BirdLists.all[idx] for idx in predictions_3[BirdLists.all].values.argmax(axis=1)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"guesses_3 = predictions_3[predictions_3.birdmax > 0.5].groupby(['audio_id', 'bird_guess'])[['site']].count().reset_index()\nguesses_3","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"site_3_birds = []\nfor audio_id, df in guesses_3.groupby('audio_id'):\n    site_3_birds.append([audio_id, ' '.join(df.bird_guess.values)])\nsite_3_birds = pd.DataFrame(site_3_birds, columns=['audio_id', 'bird_guesses'])    \nsite_3_birds = site_3_birds.merge(df_test[['audio_id', 'row_id']], how='left', on='audio_id')\nsite_3_birds","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Submission"},{"metadata":{"trusted":true},"cell_type":"code","source":"guesses = pd.concat([\n    site12_birds[['row_id', 'bird_guesses']],\n    site_3_birds[['row_id', 'bird_guesses']]\n])\nguesses.shape\nguesses.head()\nguesses.tail()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission = pd.read_csv('/kaggle/input/birdsong-recognition/sample_submission.csv')\nsubmission.head()\nsubmission.shape\n\n\nsubmission = submission.merge(guesses[['row_id', 'bird_guesses']], how='left', on=['row_id'])\n\n\nidx = (submission.bird_guesses.str.len() > 1)\nsubmission.loc[idx, 'birds'] = submission.loc[idx, 'bird_guesses']\nsubmission","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission[['row_id', 'birds']].to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"end = datetime.now()\nprint(end, (end - nb_start).seconds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}