Update tasks/audio.py

tasks/audio.py

@@ -9,8 +9,6 @@ import torch.nn.functional as F
 from torch.utils.data import DataLoader, TensorDataset
 from torchaudio import transforms
 from torchvision import models
-
-from .model import BlazeFaceModel, QuantizedBlazeFaceModel
 from .utils.evaluation import AudioEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
@@ -87,26 +85,9 @@ async def evaluate_audio(request: AudioEvaluationRequest):
         pin_memory=True,
         num_workers=4 
     )
-    
-    # Example Usage
-    model_settings = {
-        'spectrogram_length': 64,
-        'dct_coefficient_count': 481,
-        'label_count': 2
-    }    
-    model_fp32 = BlazeFaceModel(input_channels=1, label_count=model_settings['label_count'], use_double_block=False, activation='relu', use_optional_block=False)
-    quantized_model_path = "./qat_int8_blazeface_model.pth"
-
-    
-    int8_model = QuantizedBlazeFaceModel(model_fp32)
-    int8_model.eval()
-    
-    # Load the state dictionary
-    int8_model.load_state_dict(torch.load(quantized_model_path, map_location=torch.device('cpu'), weights_only=True))
-    int8_model.eval()  
 
-    #model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
-    
+    scripted_model = torch.jit.load("./optimized_qat_blazeface_model.pt", map_location=torch.device('cpu'))
+    scripted_model.eval()
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
@@ -118,7 +99,7 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     predictions = []
     with torch.inference_mode():
         for data, target in test_loader:
-            outputs = int8_model(data)
+            outputs = scripted_model(data)
             _, predicted = torch.max(outputs, 1)
             predictions.extend(predicted.tolist())
     #--------------------------------------------------------------------------------------------