Synthetic Speech Detector

import os
import zipfile
import pandas as pd
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torchvision.models import resnet18, ResNet18_Weights

import sys
sys.path.append("/bohr/dataset-gdiz/v1")
from spectrogram_dataset import SpectrogramDataset

class AudioNet(nn.Module):
    def __init__(self):
        super().__init__()
        model = resnet18(
            weights=ResNet18_Weights.DEFAULT
        )  # pretrained weights on ImageNet
        model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        model.fc = nn.Linear(model.fc.in_features, 2)
        self.model = model

    def forward(self, x):
        return self.model(x)

def train_one_epoch(model, train_loader, val_loader, criterion, optimizer, device):
    model.train()
    train_loss = 0.0

    for batch in tqdm(train_loader, desc="Train"):
        x = batch["spectrogram"].to(device)
        y = batch["label"].to(device)
        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    train_loss /= len(train_loader)
    print(f"Train Loss: {train_loss:.4f}")

    model.eval()
    val_loss = 0.0

    with torch.no_grad():
        for batch in tqdm(val_loader, desc="Val Split"):
            x = batch["spectrogram"].to(device)
            y = batch["label"].to(device)
            output = model(x)
            loss = criterion(output, y)

            val_loss += loss.item()

    val_loss /= len(val_loader)
    print(f"Val Split Loss: {val_loss:.4f}")

def predict(model, loader, device):
    model.eval()
    preds = []
    with torch.no_grad():
        for batch in tqdm(loader, desc="Test"):
            x = batch["spectrogram"].to(device)
            output = model(x)
            pred = torch.argmax(output, dim=1)
            preds.extend(pred.cpu().numpy())
    return preds

def save_submission_csv(preds, save_name):
    df = pd.DataFrame(preds)
    df.to_csv(save_name, index=False, header=False)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AudioNet().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-2)
criterion = nn.CrossEntropyLoss()

full_train_set = SpectrogramDataset("/bohr/dataset-gdiz/v1/training_set/")

val_size = int(0.2 * len(full_train_set))
train_size = len(full_train_set) - val_size
train_set, val_split_set = random_split(full_train_set, [train_size, val_size])

train_loader = DataLoader(train_set, batch_size=32)
val_split_loader = DataLoader(val_split_set, batch_size=32)

train_one_epoch(model, train_loader, val_split_loader, criterion, optimizer, device)

## Obtain the validation and test sets used for the A/B leaderboard evaluation
## [They can only be obtained normally after the notebook has been submitted to the competition]
## [An error during the code-debugging stage is normal, because the data below are not disclosed to contestants and cannot be read at this stage]

if os.environ.get('ANSWER_PATH'):
    PATH = os.environ.get("ANSWER_PATH") + "/" 
else:
    print("When the baseline runs, the test set cannot be read, so errors will follow later; this is normal.")

val_set = SpectrogramDataset(PATH+"/validation_set/")
test_set = SpectrogramDataset(PATH+"/testing_set/")

val_loader = DataLoader(val_set, batch_size=32)
test_loader = DataLoader(test_set, batch_size=32)

val_preds = predict(model, val_loader, device)
test_preds = predict(model, test_loader, device)

save_submission_csv(val_preds, "submissionA.csv")
save_submission_csv(test_preds, "submissionB.csv")
with zipfile.ZipFile("submission.zip", "w") as zipf:
    zipf.write("submissionA.csv")
    zipf.write("submissionB.csv")
os.remove("submissionA.csv")
os.remove("submissionB.csv")