!pip install dataset
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
from torchvision.models import resnet34, ResNet34_Weights
import os
import torch
from torch.utils.data import Dataset
class SpectrogramDataset(Dataset):
"""
Load spectrogram data from preprocessed .pt files.
For training_set/, assumes:
dataset/training_set/
bonafide/
spoof/
For validation_set/ and testing_set/, assumes:
dataset/validation_set/ (all .pt files in this folder, no subfolders)
dataset/testing_set/ (all .pt files in this folder, no subfolders)
No label will be provided for val/test sets to prevent label leakage.
"""
def __init__(self, directory):
self.samples = []
if "training" in directory:
label_map = {"bonafide": 0, "spoof": 1}
for label_name, label in label_map.items():
label_dir = os.path.join(directory, label_name)
if not os.path.isdir(label_dir):
continue
for fname in os.listdir(label_dir):
if fname.endswith(".pt"):
self.samples.append(
{"path": os.path.join(label_dir, fname), "label": label}
)
else:
for fname in sorted(os.listdir(directory)):
if fname.endswith(".pt"):
self.samples.append({"path": os.path.join(directory, fname)})
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
item = self.samples[idx]
spec = torch.load(item["path"])
out = {"spectrogram": spec}
if "label" in item:
out["label"] = torch.tensor(item["label"], dtype=torch.long)
return out
class AudioNet(nn.Module):
def __init__(self):
super().__init__()
base = resnet34(weights=ResNet34_Weights.DEFAULT)
self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
with torch.no_grad():
self.conv1.weight = nn.Parameter(
base.conv1.weight.mean(dim=1, keepdim=True)
)
self.bn1 = base.bn1
self.relu = base.relu
self.maxpool = base.maxpool
self.layer1 = base.layer1
self.layer2 = base.layer2
self.layer3 = base.layer3
self.layer4 = base.layer4
self.avgpool = base.avgpool
self.fc = nn.Linear(base.fc.in_features, 2)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
return self.fc(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-4, weight_decay=1e-5)
criterion = nn.CrossEntropyLoss()
full_train_set = SpectrogramDataset("/bohr/training-oa3w/v2/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, shuffle=True)
val_split_loader = DataLoader(val_split_set, batch_size=32)
for _ in range(2):
train_one_epoch(
model, train_loader, val_split_loader, criterion, optimizer, device
)
if os.environ.get('DATA_PATH'):
DATA_PATH = os.environ.get("DATA_PATH")+"/"
else:
DATA_PATH = ""
val_set = SpectrogramDataset(DATA_PATH + "/validation_set")
test_set = SpectrogramDataset(DATA_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)