File size: 6,371 Bytes
97dae9c 1eebc2d 97dae9c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
"""
This script implements a deep learning pipeline for audio classification using a pre-trained MobileNetV2 model.
The pipeline includes data loading, model training, evaluation, and emissions tracking.
"""
import os
import torch
import torch.nn as nn
import torchaudio
from torch.utils.data import Dataset, DataLoader
import numpy as np
from transformers import AutoModelForImageClassification
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from tqdm import tqdm
import logging
from datasets import load_dataset
from accelerate import Accelerator
from codecarbon import EmissionsTracker
import time
class Config:
"""
Configuration class to store hyperparameters and model settings.
"""
SAMPLE_RATE = 16000
N_FFT = 800
N_MELS = 128
HOP_LENGTH = None
SIZE = (96, 96)
SCALING_DIM = (1, 2)
LEARNING_RATE = 0.0005
BATCH_SIZE = 32
NUM_WORKERS = 4
NUM_EPOCHS = 1
MODEL_NAME = "google/mobilenet_v2_0.35_96"
MODEL_PATH = "scaled_model_800_128_96x96_mobilenet_small_unscaled_submission.pth"
config = Config()
class AudioDataset(Dataset):
"""
Custom Dataset class for loading and processing audio data.
Args:
data (list): List of audio data samples.
sample_rate (int, optional): Target sample rate for audio resampling. Defaults to 16000.
audio_target_length (float, optional): Target length of audio in seconds. Defaults to 4.5.
"""
def __init__(self, data, sample_rate=16000, audio_target_length=4.5):
self.data = data
self.sample_rate = sample_rate
self.audio_target_length = audio_target_length
def __len__(self):
return len(self.data)
def __getitem__(self, index):
# 1. Cache the resampler
if not hasattr(self, '_resampler_cache'):
self._resampler_cache = {}
# 2. Get data efficiently
data_item = self.data[index]
waveform = torch.FloatTensor(data_item["audio"]["array"]) if len(data_item["audio"]["array"]) > 0 else torch.ones(36000)*1E-5
# 4. Cached resampler creation
orig_freq = waveform.shape[-1]
target_freq = self.audio_target_length * self.sample_rate
resampler_key = (orig_freq, target_freq)
if resampler_key not in self._resampler_cache:
self._resampler_cache[resampler_key] = torchaudio.transforms.Resample(
orig_freq=orig_freq,
new_freq=target_freq
)
# 5. Apply resampling and return
return self._resampler_cache[resampler_key](waveform), data_item["label"]
def collate_fn(batch):
"""
Collate function to stack inputs and labels into batches.
Args:
batch (list): List of tuples containing inputs and labels.
Returns:
tuple: Stacked inputs and labels.
"""
return torch.stack([inputs for inputs, _ in batch]), torch.tensor([label for _, label in batch])
class AudioClassifier(nn.Module):
"""
Audio classification model using a pre-trained MobileNetV2.
Args:
model_name (str): Name of the pre-trained model.
model_path (str): Path to save/load the model.
new (bool, optional): Whether to load a new model or an existing one. Defaults to True.
"""
def __init__(self, model_name, model_path, new=True):
super().__init__()
self.model = self.load_model(model_name, model_path, new)
self.num_classes = 2
self.mel_spectrogram = torchaudio.transforms.MelSpectrogram(
sample_rate=config.SAMPLE_RATE,
n_fft=config.N_FFT,
n_mels=config.N_MELS,
hop_length=config.HOP_LENGTH
)
self.amplitude_to_db = torchaudio.transforms.AmplitudeToDB()
def load_model(self, model_name, model_path, new=False):
"""
Load the pre-trained model and modify the classifier.
Args:
model_name (str): Name of the pre-trained model.
model_path (str): Path to save/load the model.
new (bool, optional): Whether to load a new model or an existing one. Defaults to False.
Returns:
nn.Module: Loaded model.
"""
model = AutoModelForImageClassification.from_pretrained(model_name)
model.classifier = torch.nn.Sequential(
nn.Linear(in_features=1280, out_features=2))
for param in model.parameters():
param.requires_grad = True
state_dict = torch.load(model_path)
model.load_state_dict(state_dict)
return model
def forward(self, waveforms):
"""
Forward pass through the model.
Args:
waveforms (torch.Tensor): Input audio waveforms.
Returns:
torch.Tensor: Model output.
"""
melspectrogram = self.mel_spectrogram(waveforms)
melspectrogram = nn.functional.interpolate(melspectrogram.unsqueeze(1),
size=config.SIZE,
mode="bilinear",
align_corners=False).squeeze(1)
db_melspectrogram = self.amplitude_to_db(melspectrogram)
delta = torchaudio.functional.compute_deltas(melspectrogram)
x = torch.stack([melspectrogram, db_melspectrogram, delta], dim=1)
return self.model(x)
class Evaluator:
def __init__(self, model, dataloader, device):
self.model = model
self.dataloader = dataloader
self.device = device
@torch.no_grad()
def evaluate(self):
self.model.eval()
all_predictions = []
all_labels = []
idx = 0
for waveforms, labels in self.dataloader:
waveforms = waveforms.to(self.device)
outputs = self.model(waveforms).logits
predictions = torch.argmax(outputs, dim=1)
all_predictions.extend(predictions.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
idx += 1
if idx % 10 == 0:
torch.cuda.empty_cache()
all_predictions = np.array(all_predictions)
all_labels = np.array(all_labels)
# return self.compute_metrics(all_predictions, all_labels)
return all_predictions
|