gpu to cpu

get_difficulty.py

@@ -32,18 +32,16 @@ def get_cqt_from_mp3(mp3_path):
     log_cqt = log_cqt.T  # shape (T, 88)
     log_cqt = downsample_log_cqt(log_cqt, target_fs=5)
     cqt_tensor = torch.tensor(log_cqt, dtype=torch.float32).unsqueeze(0).unsqueeze(0).cpu()
-    # pdb.set_trace()
     print(f"cqt shape: {log_cqt.shape}")
     return cqt_tensor
 
 def get_pianoroll_from_mp3(mp3_path):
     audio, _ = load_audio(mp3_path, sr=sample_rate, mono=True)
-    transcriptor = PianoTranscription(device='cpu')
+    transcriptor = PianoTranscription(device="cuda" if torch.cuda.is_available() else "cpu")
     midi_path = "temp.mid"
     transcriptor.transcribe(audio, midi_path)
     midi_data = pretty_midi.PrettyMIDI(midi_path)
 
-    # Create pianoroll and onset matrix
     fs = 5  # original frames per second
     piano_roll = midi_data.get_piano_roll(fs=fs)[21:109].T  # shape: (T, 88)
     piano_roll = piano_roll / 127
@@ -64,6 +62,8 @@ def get_pianoroll_from_mp3(mp3_path):
     return out_tensor.transpose(2, 3)
 
 def predict_difficulty(mp3_path, model_name, rep):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     if "only_cqt" in rep:
         only_cqt, only_pr = True, False
         rep_clean = "multimodal5"
@@ -74,18 +74,17 @@ def predict_difficulty(mp3_path, model_name, rep):
         only_cqt = only_pr = False
         rep_clean = rep
 
-    model = AudioModel(num_classes=11, rep=rep_clean, modality_dropout=False, only_cqt=only_cqt, only_pr=only_pr)
-    checkpoint = [torch.load(f"models/{model_name}/checkpoint_{i}.pth", map_location="cpu", weights_only=False)
+    model = AudioModel(num_classes=11, rep=rep_clean, modality_dropout=False, only_cqt=only_cqt, only_pr=only_pr).to(device)
+    checkpoint = [torch.load(f"models/{model_name}/checkpoint_{i}.pth", map_location=device, weights_only=False)
                   for i in range(5)]
 
-
     if rep == "cqt5":
-        inp_data = get_cqt_from_mp3(mp3_path)
+        inp_data = get_cqt_from_mp3(mp3_path).to(device)
     elif rep == "pianoroll5":
-        inp_data = get_pianoroll_from_mp3(mp3_path)
+        inp_data = get_pianoroll_from_mp3(mp3_path).to(device)
     elif rep_clean == "multimodal5":
-        x1 = get_pianoroll_from_mp3(mp3_path)
-        x2 = get_cqt_from_mp3(mp3_path)
+        x1 = get_pianoroll_from_mp3(mp3_path).to(device)
+        x2 = get_cqt_from_mp3(mp3_path).to(device)
         inp_data = [x1, x2]
     else:
         raise ValueError(f"Representation {rep} not supported")
@@ -93,23 +92,16 @@ def predict_difficulty(mp3_path, model_name, rep):
     preds = []
     for cheks in checkpoint:
         model.load_state_dict(cheks["model_state_dict"])
-        model = model.cpu().eval()
+        model.eval()
         with torch.inference_mode():
             logits = model(inp_data, None)
             pred = prediction2label(logits).item()
             preds.append(pred)
 
     return mean(preds)
-    # return preds
 
 if __name__ == "__main__":
     mp3_path = "yt_audio.mp3"
-    model_name = ""
-    # pred_cqt = predict_difficulty(mp3_path, model_name="audio_midi_cqt5_ps_v5", rep="cqt5")
-    # print(f"Predicción dificultad CQT: {pred_cqt}")
-
-    # pred_pr = predict_difficulty(mp3_path, model_name="audio_midi_pianoroll_ps_5_v4", rep="pianoroll5")
-    # print(f"Predicción dificultad PR: {pred_pr}")
-
-    pred_multi = predict_difficulty(mp3_path, model_name="audio_midi_multi_ps_v5", rep="multimodal5")
-    print(f"Predicción dificultad multimodal: {pred_multi}")
+    model_name = "audio_midi_multi_ps_v5"
+    pred_multi = predict_difficulty(mp3_path, model_name=model_name, rep="multimodal5")
+    print(f"Predicción dificultad multimodal: {pred_multi}")

model.py

@@ -212,127 +212,6 @@ class AudioModel(nn.Module):
         return x
 
 
-def get_mse_macro(y_true, y_pred):
-    mse_each_class = []
-    for true_class in set(y_true):
-        tt, pp = zip(*[[tt, pp] for tt, pp in zip(y_true, y_pred) if tt == true_class])
-        mse_each_class.append(mean_squared_error(y_true=tt, y_pred=pp))
-    return mean(mse_each_class)
-
-
-def get_cqt(rep, k):
-    inp_data = utils.load_binary(f"../videos_download/{rep}/{k}.bin")
-    inp_data = torch.tensor(inp_data, dtype=torch.float32).cpu()
-    inp_data = inp_data.unsqueeze(0).unsqueeze(0).transpose(2, 3)
-    return inp_data
-
-
-def get_pianoroll(rep, k):
-    inp_pr = utils.load_binary(f"../videos_download/{rep}/{k}.bin")
-    inp_on = utils.load_binary(f"../videos_download/{rep}/{k}_onset.bin")
-    inp_pr = torch.from_numpy(inp_pr).float().cpu()
-    inp_on = torch.from_numpy(inp_on).float().cpu()
-    inp_data = torch.stack([inp_pr, inp_on], dim=1)
-    inp_data = inp_data.unsqueeze(0).permute(0, 1, 2, 3)
-    return inp_data
-
-def compute_model_basic(model_name, rep, modality_dropout, only_cqt=False, only_pr=False):
-    seed = 42
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed(seed)
-    data = utils.load_json("../videos_download/split_audio.json")
-    mse, acc = [], []
-    predictions = []
-    if only_cqt:
-        cache_name = model_name + "_cqt"
-    elif only_pr:
-        cache_name = model_name + "_pr"
-    else:
-        cache_name = model_name
-    if not os.path.exists(f"cache/{cache_name}.json"):
-        for split in range(5):
-            #load_model
-            model = AudioModel(11, rep, modality_dropout, only_cqt, only_pr)
-            checkpoint = torch.load(f"models/{model_name}/checkpoint_{split}.pth",  map_location='cpu')
-            # print(checkpoint["epoch"])
-            # print(checkpoint.keys())
-
-            model.load_state_dict(checkpoint['model_state_dict'])
-            model = model.cpu()
-            pred_labels, true_labels = [], []
-            predictions_split = {}
-            model.eval()
-            with torch.inference_mode():
-                for k, ps in data[str(split)]["test"].items():
-                    # computar el modelo
-                    if "cqt" in rep:
-                        inp_data = get_cqt(rep, k)
-                    elif "pianoroll" in rep:
-                        inp_data = get_pianoroll(rep, k)
-                    elif rep == "multimodal5":
-                        x1 = get_pianoroll("pianoroll5", k)
-                        x2 = get_cqt("cqt5", k)[:, :, :x1.shape[2]]
-                        inp_data = [x1, x2]
-                    log_prob = model(inp_data, None)
-                    pred = prediction2label(log_prob).cpu().tolist()[0]
-                    print(k, ps, pred)
-                    predictions_split[k] = {
-                        "true": ps,
-                        "pred": pred
-                    }
-
-                    true_labels.append(ps)
-                    pred_labels.append(pred)
-
-            predictions.append(predictions_split)
-            mse.append(get_mse_macro(true_labels, pred_labels))
-            acc.append(balanced_accuracy_score(true_labels, pred_labels))
-        # with one decimal
-        print(f"mse: {mean(mse):.1f}({stdev(mse):.1f})", end=" ")
-        print(f"acc: {mean(acc)*100:.1f}({stdev(acc)*100:.1f})")
-        utils.save_json({
-            "mse": mse,
-            "acc": acc,
-            "predictions": predictions
-        }, f"cache/{cache_name}.json")
-    else:
-        data = utils.load_json(f"cache/{cache_name}.json")
-        tau_c, mse, acc = [], [], []
-        for i in range(5):
-            pred, true = [], []
-            for k, dd in data["predictions"][i].items():
-                pred.append(dd["pred"])
-                true.append(dd["true"])
-            tau_c.append(kendalltau(x=true, y=pred).statistic)
-            mse.append(get_mse_macro(true, pred))
-            acc.append(balanced_accuracy_score(true, pred))
-        print(model_name, end="// ")
-        print(f"& {mean(mse):.2f}({stdev(mse):.2f})", end=" ")
-        print(f"& {mean(acc) * 100:.1f}({stdev(acc) * 100:.2f})", end=" ")
-        print(f"& {mean(tau_c):.3f}({stdev(tau_c):.3f})")
-
-
-def compute_ensemble(truncate=False):
-    round_func = lambda x: math.ceil(x) if truncate else math.floor(x)
-    data_pr = utils.load_json(f"cache/audio_midi_cqt5_ps_v5.json")
-    data_cqt = utils.load_json(f"cache/audio_midi_pianoroll_ps_5_v4.json")
-    tau_c, mse, acc = [], [], []
-    for i in range(5):
-        pred, true = [], []
-        for k, dd in data_pr["predictions"][i].items():
-            cqt_pred = data_cqt["predictions"][i][k]
-            pred.append(round_func((dd["pred"] + cqt_pred["pred"])/2))
-            true.append(dd["true"])
-        tau_c.append(kendalltau(x=true, y=pred).statistic)
-        mse.append(get_mse_macro(true, pred))
-        acc.append(balanced_accuracy_score(true, pred))
-    print("ensemble", end="// ")
-    print(f"& {mean(mse):.2f}({stdev(mse):.2f})", end=" ")
-    print(f"& {mean(acc) * 100:.1f}({stdev(acc) * 100:.2f})", end=" ")
-    print(f"& {mean(tau_c):.3f}({stdev(tau_c):.3f})")
-
 
 def load_json(name_file):
     with open(name_file, 'r') as fp: