Update tasks/audio.py

tasks/audio.py

@@ -66,25 +66,22 @@ async def evaluate_audio(request: AudioEvaluationRequest):
             _waveform = _resampler(_waveform)
         return _waveform
 
-    resized_waveforms = [
-        resize_audio(torch.tensor(sample['audio']['array'], dtype=torch.float32).unsqueeze(0), target_length=72000)
-        for sample in test_dataset
-    ]
+    def preprocess_audio(sample):
+        waveform = torch.tensor(sample['audio']['array'], dtype=torch.float32).unsqueeze(0)
+        resized_waveform = resize_audio(waveform, target_length=72000)
+        return amplitude_to_db(mel_transform(resampler(resized_waveform)))
 
-    waveforms, labels = [], []
-    for waveform, label in zip(resized_waveforms, true_labels):
-        waveforms.append(amplitude_to_db(mel_transform(resampler(waveform))))
-        labels.append(label)
+    waveforms = [preprocess_audio(sample) for sample in test_dataset]
+    labels = torch.tensor(true_labels)
 
     waveforms = torch.stack(waveforms)
-    labels = torch.tensor(labels)
 
     test_loader = DataLoader(
         TensorDataset(waveforms, labels),
         batch_size=128,
         shuffle=False,
         pin_memory=True,
-        num_workers=2
+        num_workers=4
     )
 
     scripted_model = torch.jit.load("./optimized_qat_blazeface_model.pt", map_location=torch.device('cpu'))
@@ -101,8 +98,7 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     predictions = []
     with torch.no_grad():
         #with autocast():
-        #with torch.amp.autocast(device_type='cpu'):
-        with torch.autocast(device_type='cpu'):
+        with torch.amp.autocast(device_type='cpu'):
             for data, target in test_loader:
                 outputs = scripted_model(data)
                 _, predicted = torch.max(outputs, 1)