<jupyter_start><jupyter_text>Denoising mel-spectrograms with a residual dense network  Setup  Environment<jupyter_code>!pip install -q ipython_secrets gsheet-keyring comet_ml
%%writefile setup.sh

export CUDA_HOME=/usr/local/cuda-10.1
[[ -d ./apex ]] || git clone https://github.com/NVIDIA/apex && pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./apex
!sh setup.sh<jupyter_output>setup.sh: 3: setup.sh: [[: not found
fatal: destination path 'apex' already exists and is not an empty directory.<jupyter_text>Variables<jupyter_code>from ipython_secrets import get_secret

CODEBASE = "https://github.com/sdll/audio-denoising"

PROJECT = "Arbeit"
COMET_ML_API_KEY = get_secret("comet-{}".format(PROJECT))

GDRIVE_MOUNT_POINT = "/content/drive"
TRAIN_DATASET = "dataset/train"
VAL_DATASET = "dataset/val"
WORKSPACE= "sdll"

from google.colab import drive
import os
drive.mount(GDRIVE_MOUNT_POINT)
GOOGLE_DRIVE_ROOT = GDRIVE_MOUNT_POINT + "/" + list(filter(lambda x: x[0] != '.', os.listdir(GDRIVE_MOUNT_POINT)))[0]<jupyter_output>Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount("/content/drive", force_remount=True).<jupyter_text>Data<jupyter_code>!git clone $CODEBASE
%cd audio-denoising
!mkdir -p dataset

![[ -f ./dataset/train.zip ]] || wget https://www.dropbox.com/s/n6nhp5e231rl0b5/train.zip?raw=1 -O dataset/train.zip
![[ -f ./dataset/test.zip ]] || wget https://www.dropbox.com/s/nt4q2n0esiboc1i/val.zip?raw=1 -O dataset/val.zip
%cd dataset
![[ -d train ]] || unzip -qq train.zip
![[ -d val ]] || unzip -qq val.zip
%cd ..
![[ -d pytorch_ssim ]] || (git clone https://github.com/Po-Hsun-Su/pytorch-ssim && mv pytorch-ssim/pytorch_ssim . && rm -rf pytorch-ssim)<jupyter_output><empty_output><jupyter_text>Imports<jupyter_code>%cd /content/audio-denoising/
from comet_ml import Experiment

import argparse
from timeit import default_timer as timer

import numpy as np
import seaborn as sns
import torch
from matplotlib import pyplot as plt
from torch import nn, optim
from tqdm import tqdm as tqdm_base

from audio_denoising.data.loader import load
from audio_denoising.model.rdn import ResidualDenseNetwork as Model
from pytorch_ssim import ssim


def tqdm(*args, **kwargs):
    if hasattr(tqdm_base, "_instances"):
        for instance in list(tqdm_base._instances):
            tqdm_base._decr_instances(instance)
    return tqdm_base(*args, **kwargs)
from apex import amp, optimizers
sns.set()<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>def get_arg_parser():
    parser = argparse.ArgumentParser()

    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--growth-rate", type=int, default=16)
    parser.add_argument("--kernel-size", type=int, default=3)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--num-blocks", type=int, default=20)
    parser.add_argument("--num-channels", type=int, default=1)
    parser.add_argument("--num-epochs", type=int, default=80)
    parser.add_argument("--num-features", type=int, default=16)
    parser.add_argument("--num-layers", type=int, default=6)
    parser.add_argument("--seed", type=int, default=42)

    return parser


def get_criterion():
    return nn.MSELoss()


def get_optimizer(model, lr=1e-3):
    return optimizers.FusedAdam(model.parameters(), lr)
    # return optim.SGD(model.parameters(), lr)


def psnr(prediction, target, max_pixel=255.0):
    return 10.0 * ((max_pixel ** 2) / ((prediction - target) ** 2).mean()).log10()


def train(
    experiment,
    loader,
    model,
    criterion,
    optimizer,
    args,
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
    scheduler=None,
    validator=None,  # (validate, val_loader)
    verbose=True,
):
    experiment.log_parameters(vars(args))

    np.random.seed(args.seed)

    # save the model weights with the best psnr
    if validator:
        best_psnr = 0.0

    for epoch in tqdm(range(args.num_epochs), desc="Epoch", unit="epochs"):
        with experiment.train():
            model.train()
            train_psnr = []
            train_ssim = []

            for clean_image, noisy_image in tqdm(
                loader, desc="Train images", unit="batches"
            ):
                image = noisy_image.to(device, dtype=torch.float)
                gt_image = clean_image.to(device, dtype=torch.float)
                noise = image - gt_image

                prediction = model(image)
                cleaned_image = image - prediction

                loss = criterion(prediction, noise)

                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()

                optimizer.step()
                optimizer.zero_grad()

                current_psnr = psnr(cleaned_image, gt_image).data.item()
                current_ssim = ssim(cleaned_image, gt_image).data.item()
                train_psnr.append(current_psnr)
                train_ssim.append(current_ssim)
                experiment.log_metric("psnr", current_psnr)
                experiment.log_metric("ssim", current_ssim)
                experiment.log_metric("loss", scaled_loss.data.item())

            experiment.log_metric("mean_psnr", np.mean(train_psnr))
            experiment.log_metric("mean_ssim", np.mean(train_ssim))
        if validator:
            validate, val_loader = validator
            test_psnr, _, _ = validate(experiment, val_loader, model, device, verbose)
            if scheduler:
                scheduler.step(test_psnr)
            if test_psnr > best_psnr:
                best_psnr = test_psnr
                filename_pth = GOOGLE_DRIVE_ROOT + '/audio_denoising_psnr_{:.4f}_epoch_{}_D_{}_C_{}_G_{}_G0_{}.pth'.format(
                    test_psnr, epoch, args.num_blocks, args.num_layers, args.growth_rate, args.num_features
                )
                torch.save(model.state_dict(), filename_pth)
    return model


def validate(
    experiment,
    loader,
    model,
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
    verbose=True,
):
    with experiment.test():
        model.eval()
        test_psnr = []
        test_ssim = []
        test_prediction_times = []
        for clean_image, noisy_image in tqdm(loader, desc="Val images", unit="batches"):
            image = noisy_image.to(device, dtype=torch.float)
            gt_image = clean_image.to(device, dtype=torch.float)
            start = timer()
            prediction = model(image)
            end = timer()
            cleaned_image = image - prediction

            prediction_time = end - start
            test_prediction_times.append(prediction_time)
            experiment.log_metric("prediction_time", prediction_time)

            current_psnr = psnr(cleaned_image, gt_image).data.item()
            current_ssim = ssim(cleaned_image, gt_image).data.item()

            test_psnr.append(current_psnr)
            test_ssim.append(current_ssim)

        test_psnr = np.mean(test_psnr)
        test_ssim = np.mean(test_ssim)
        test_prediction_time = np.mean(test_prediction_times)

        experiment.log_metric("mean_psnr", test_psnr)
        experiment.log_metric("mean_ssim", test_ssim)
        experiment.log_metric("mean_prediction_time", test_prediction_time)

    if verbose:
        print(
            "\nMean Test PSNR: {:.2f}\nMean Test SSIM: {:.2f}\nMean Prediction Time: {:.2f}".format(
                test_psnr, test_ssim, test_prediction_time
            )
        )
    return test_psnr, test_ssim, test_prediction_time
experiment = Experiment(
    api_key=COMET_ML_API_KEY,
    project_name=PROJECT,
    workspace=WORKSPACE,
    auto_output_logging=None,
)
args = get_arg_parser().parse_args(args=[])

train_loader = load(TRAIN_DATASET, batch_size=args.batch_size)
val_loader = load(VAL_DATASET, batch_size=1)

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Model(args).to(device)
optimizer = get_optimizer(model, args.lr)
model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
criterion = get_criterion()
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.75, patience=0, verbose=True)
train(
    experiment,
    train_loader,
    model,
    criterion,
    optimizer,
    args,
    device,
    scheduler=scheduler,
    validator=(validate, val_loader)
)<jupyter_output>Train images: 100%|██████████| 1500/1500 [17:53<00:00,  1.40batches/s]
Val images: 100%|██████████| 2000/2000 [01:38<00:00, 20.40batches/s]
Epoch:   1%|▏         | 1/80 [19:32<25:43:24, 1172.21s/epochs]<jupyter_text>The Comet logs of this round of training can be found [here](https://www.comet.ml/sdll/arbeit/a40a28c88443436984eedc52913c3a09).<jupyter_code>filename_pth = GOOGLE_DRIVE_ROOT + '/audio_denoising_psnr_{:.4f}_epoch_{}.pth'.format(
    59.71, 35 
)
torch.save(model.state_dict(), filename_pth)<jupyter_output><empty_output><jupyter_text>Validation<jupyter_code># ckpt = GOOGLE_DRIVE_ROOT + '/audio_denoising_psnr_59.7472_epoch_6.pth'
ckpt = GOOGLE_DRIVE_ROOT + "/audio_denoising_psnr_65.1736_epoch_15_D_20_C_6_G_16_G0_16.pth"
model.load_state_dict(torch.load(ckpt))
test_psnr, test_ssim, test_prediction_time = validate(experiment, val_loader, model, verbose=True)

%matplotlib inline
from IPython.display import clear_output, display
import sys
import time
for clean_mel, noisy_mel in val_loader: 
    time.sleep(.25)
    clean_mel = clean_mel.to("cpu", dtype=torch.float)
    noisy_mel = noisy_mel.to("cpu", dtype=torch.float)
    fig = plt.figure(figsize=(17, 8), dpi=100)
    ax1 = fig.add_subplot(121)
    ax2 = fig.add_subplot(122)
    ax1.imshow(clean_mel.squeeze(), interpolation='nearest', aspect='auto')
    ax1.set_title('Clean image')
    ax2.imshow((noisy_mel - model(noisy_mel.to(device)).data.to("cpu")).squeeze(),
               interpolation='nearest', aspect='auto')
    ax2.set_title('Cleaned image')
    display(fig)
    clear_output(wait=True)
    sys.stdout.flush()
    break<jupyter_output><empty_output><jupyter_text>SSIM Inference Analysis<jupyter_code>from audio_denoising.data.loader import SpectogramDataset
clean_dataset = SpectogramDataset(VAL_DATASET + "/clean")
noisy_dataset = SpectogramDataset(VAL_DATASET + "/noisy")
def compute_ssims(dataset):
    ssims = []
    for file_idx in tqdm(range(len(dataset)), unit="files"):
        img = dataset[file_idx].unsqueeze(0).to(device, dtype=torch.float)
        noise = model(img).to("cpu")
        img = img.to("cpu")
        ssims.append(ssim(img - noise, img).data.item())
    return ssims
clean_ssims = compute_ssims(clean_dataset)
plt.figure(dpi=100)
sns.distplot(clean_ssims)
plt.show()
noisy_ssims = compute_ssims(noisy_dataset)
plt.figure(dpi=100)
sns.distplot(noisy_ssims)
plt.show()
clean_ssims = np.asarray(clean_ssims)
noisy_ssims = np.asarray(noisy_ssims)
threshold = 0.77
true_clean = np.sum(clean_ssims > threshold) / len(clean_ssims)
true_noisy = np.sum(noisy_ssims < threshold) / len(noisy_ssims)
print("True clean ratio: {}".format(true_clean))
print("True noisy ratio: {}".format(true_noisy))<jupyter_output>True clean ratio: 0.95
True noisy ratio: 0.9425<jupyter_text>Scratchpad<jupyter_code>dataset = load(TRAIN_DATASET, batch_size=1)
heights = []
for clean_mel, noisy_mel in dataset:
    heights.append(clean_mel.shape[2]) 
plt.figure(dpi=120)
plt.title("Distribution of Train MEL Heights")
plt.xlabel("Height, px")
plt.ylabel("Frequency")
sns.distplot(heights)
plt.show()
min(heights), max(heights)
%cd /content/audio-denoising/
!git pull
%%capture
args = get_arg_parser().parse_args(args=[])

train_loader = load(TRAIN_DATASET, batch_size=args.batch_size)
val_loader = load(VAL_DATASET, batch_size=1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

experiment = Experiment(
    api_key=COMET_ML_API_KEY,
    project_name=PROJECT,
    workspace=WORKSPACE,
    auto_output_logging=None,
)

grid_search = {
    "num_blocks": np.arange(3, 12, 3),
    "num_layers": np.arange(6, 18, 3),
    "num_features": np.arange(8, 32, 8),
    "growth_rate": np.arange(8, 32, 8),
}

best = {
    "num_blocks": None,
    "num_layers": None,
    "num_features": None,
    "growth_rate": None,
}

args.num_epochs = 2

best_psnr = 0.0
for num_blocks in grid_search["num_blocks"]:
    for num_layers in grid_search["num_layers"]:
        for num_features in grid_search["num_features"]:
            for growth_rate in grid_search["growth_rate"]:
                setattr(args, "num_blocks", num_blocks)
                setattr(args, "num_layers", num_layers)
                setattr(args, "num_features", num_features)
                setattr(args, "growth_rate", growth_rate)
                model = Model(args).to(device)
                optimizer = get_optimizer(model, args.lr)
                model, optimizer = amp.initialize(model, optimizer, opt_level="O3")
                criterion = get_criterion()
                scheduler = optim.lr_scheduler.ReduceLROnPlateau(
                    optimizer, mode="max", factor=0.75, patience=0, verbose=True
                )
                train(
                    experiment,
                    [next(iter(train_loader))],
                    model,
                    criterion,
                    optimizer,
                    args,
                    device,
                    scheduler=scheduler,
                )
                test_psnr, _, _ = validate(
                    experiment, [next(iter(val_loader))], model, device, verbose=False
                )
                if test_psnr > best_psnr:
                    best_psnr = test_psnr
                    best["num_blocks"] = num_blocks
                    best["num_layers"] = num_layers
                    best["num_features"] = num_features
                    best["growth_rate"] = growth_rate
                del model, optimizer, criterion, scheduler
                torch.cuda.empty_cache()
experiment.end()

print(best)
while True:
    pass
next(iter(val_loader))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12440564899581293, 'fraction_numerical': 0.01752891916826343, 'mean_word_length': 3.217898832684825, 'pattern_counts': {'":': 8, '<': 37, '<?xml version=': 0, '>': 36, 'https://': 6, 'lorem ipsum': 0, 'www.': 3, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '495151', 'license': 'MIT', 'n_tokens_mistral': 5089, 'n_tokens_neox': 4835, 'n_words': 855, 'path': 'demo.ipynb', 'repo_name': 'sdll/audio-denoising', 'size': '14091'}