remove pwvt

tasks/run_inr.py

@@ -58,11 +58,11 @@ local_rank = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 login(api_key)
 dataset = load_dataset("rfcx/frugalai", streaming=True)
 
-train_ds = SplitDataset(FFTDataset(dataset["train"]), is_train=True)
+train_ds = SplitDataset(AudioINRDataset(FFTDataset(dataset["train"])), is_train=True)
 
 train_dl = DataLoader(train_ds, batch_size=data_args.batch_size)
 
-val_ds = SplitDataset(FFTDataset(dataset["train"]), is_train=False)
+val_ds = SplitDataset(AudioINRDataset(FFTDataset(dataset["train"])), is_train=False)
 
 val_dl = DataLoader(val_ds, batch_size=data_args.batch_size)
 
@@ -110,34 +110,34 @@ test_dl = DataLoader(test_ds, batch_size=data_args.batch_size)
 loss_fn = torch.nn.BCEWithLogitsLoss()
 inr_criterion = torch.nn.MSELoss()
 
-# for i, batch in enumerate(train_ds):
-#     coords, fft, audio = batch['audio']['coords'], batch['audio']['fft_mag'], batch['audio']['array']
-#     coords = coords.to(local_rank)
-#     fft = fft.to(local_rank)
-#     audio = audio.to(local_rank)
-#     values = torch.cat((audio.unsqueeze(-1), fft.unsqueeze(-1)), dim=-1)
-#     # model = INR(hidden_features=128, n_layers=3,
-#     #             in_features=1,
-#     #             out_features=1).to(local_rank)
-#     model = FasterKAN(**kan_args.get_dict()).to(local_rank)
-#     optimizer = torch.optim.Adam([{'params': model.parameters()}], lr=1e-3)
-#     pbar = tqdm(range(200))
-#     losses = []
-#     print(coords.shape)
-#     for t in pbar:
-#         optimizer.zero_grad()
-#         pred_values = model(coords.to(local_rank)).float()
-#         loss = inr_criterion(pred_values, values)
-#         loss.backward()
-#         optimizer.step()
-#         pbar.set_description(f'loss: {loss.item()}')
-#         losses.append(loss.item())
-#     state_dict = model.state_dict()
-#     torch.save(state_dict, 'test')
-#     # print(f'Sample {i+offset} label {label} saved in {inr_path}')
-#     plot_results(1, i, fft, losses, pred_values)
+for i, batch in enumerate(train_ds):
+    coords, fft, audio = batch['audio']['coords'], batch['audio']['fft_mag'], batch['audio']['array']
+    coords = coords.to(local_rank)
+    fft = fft.to(local_rank)
+    audio = audio.to(local_rank)
+    # values = torch.cat((audio.unsqueeze(-1), fft.unsqueeze(-1)), dim=-1)
+    model = INR(hidden_features=128, n_layers=4,
+                in_features=1,
+                out_features=1).to(local_rank)
+    # model = FasterKAN(layers_hidden=[1,16,16,1]).to(local_rank)
+    optimizer = torch.optim.Adam([{'params': model.parameters()}], lr=1e-3)
+    pbar = tqdm(range(200))
+    losses = []
+    print(coords.shape)
+    for t in pbar:
+        optimizer.zero_grad()
+        pred_values = model(coords.to(local_rank)).float()
+        loss = inr_criterion(pred_values, fft)
+        loss.backward()
+        optimizer.step()
+        pbar.set_description(f'loss: {loss.item()}')
+        losses.append(loss.item())
+    state_dict = model.state_dict()
+    torch.save(state_dict, 'test')
+    # print(f'Sample {i+offset} label {label} saved in {inr_path}')
+    plot_results(1, i, fft, losses, pred_values)
 # #
-# exit()
+exit()
 
 
 # missing, unexpected = model.load_state_dict(torch.load(model_args.checkpoint_path))

tasks/utils/transforms.py

@@ -2,7 +2,7 @@ import numpy as np
 import librosa
 import torch
 import torch.nn as nn
-import pywt
+# import pywt
 from scipy import signal
 
 
@@ -59,31 +59,31 @@ def compute_cwt_power_spectrum(audio, sample_rate, num_freqs=128, f_min=20, f_ma
 
     return power_spectrum_tensor
 
-def compute_wavelet_transform(audio, wavelet, decompos_level):
-    """Compute wavelet decomposition of the audio signal."""
-    # Convert to numpy and ensure 1D
-    audio_np = audio.cpu().numpy()
-
-    # Perform wavelet decomposition
-    coeffs = pywt.wavedec(audio_np, wavelet, level=decompos_level)
-
-    # Stack coefficients into a 2D array
-    # First, pad all coefficient arrays to the same length
-    max_len = max(len(c) for c in coeffs)
-    padded_coeffs = []
-    for coeff in coeffs:
-        pad_len = max_len - len(coeff)
-        if pad_len > 0:
-            padded_coeff = np.pad(coeff, (0, pad_len), mode='constant')
-        else:
-            padded_coeff = coeff
-        padded_coeffs.append(padded_coeff)
-
-    # Stack into 2D array where each row is a different scale
-    wavelet_features = np.stack(padded_coeffs)
-
-    # Convert to tensor
-    return torch.FloatTensor(wavelet_features)
+# def compute_wavelet_transform(audio, wavelet, decompos_level):
+#     """Compute wavelet decomposition of the audio signal."""
+#     # Convert to numpy and ensure 1D
+#     audio_np = audio.cpu().numpy()
+#
+#     # Perform wavelet decomposition
+#     coeffs = pywt.wavedec(audio_np, wavelet, level=decompos_level)
+#
+#     # Stack coefficients into a 2D array
+#     # First, pad all coefficient arrays to the same length
+#     max_len = max(len(c) for c in coeffs)
+#     padded_coeffs = []
+#     for coeff in coeffs:
+#         pad_len = max_len - len(coeff)
+#         if pad_len > 0:
+#             padded_coeff = np.pad(coeff, (0, pad_len), mode='constant')
+#         else:
+#             padded_coeff = coeff
+#         padded_coeffs.append(padded_coeff)
+#
+#     # Stack into 2D array where each row is a different scale
+#     wavelet_features = np.stack(padded_coeffs)
+#
+#     # Convert to tensor
+#     return torch.FloatTensor(wavelet_features)
 
 
 def compute_melspectrogram(audio, sample_rate):