122 lines
6.0 KiB
Python
122 lines
6.0 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
import wandb
|
|
import random
|
|
import os
|
|
import pickle
|
|
from data_processing import delta_encode, delta_decode, save_wav
|
|
from utils import visualize_prediction, plot_delta_distribution
|
|
from bitstream import ArithmeticEncoder
|
|
|
|
def pad_sequence(sequence, max_length):
|
|
padded_seq = np.zeros((max_length, *sequence.shape[1:]))
|
|
padded_seq[:sequence.shape[0], ...] = sequence
|
|
return padded_seq
|
|
|
|
def evaluate_model(model, data, use_delta_encoding, encoder, sample_rate=19531, epoch=0):
|
|
compression_ratios = []
|
|
identical_count = 0
|
|
all_deltas = []
|
|
|
|
for i, file_data in enumerate(data):
|
|
file_data = torch.tensor(file_data, dtype=torch.float32).unsqueeze(1).to(model.device)
|
|
encoded_data = model.encode(file_data.squeeze(1).cpu().numpy())
|
|
encoder.build_model(encoded_data)
|
|
compressed_data = encoder.encode(encoded_data)
|
|
decompressed_data = encoder.decode(compressed_data, len(encoded_data))
|
|
|
|
if use_delta_encoding:
|
|
decompressed_data = delta_decode(decompressed_data)
|
|
|
|
# Ensure the lengths match
|
|
min_length = min(len(file_data), len(decompressed_data))
|
|
file_data = file_data[:min_length]
|
|
decompressed_data = decompressed_data[:min_length]
|
|
|
|
identical = np.allclose(file_data.cpu().numpy(), decompressed_data, atol=1e-5)
|
|
if identical:
|
|
identical_count += 1
|
|
|
|
compression_ratio = len(file_data) / len(compressed_data)
|
|
compression_ratios.append(compression_ratio)
|
|
|
|
predicted_data = model(torch.tensor(encoded_data, dtype=torch.float32).unsqueeze(1).to(model.device)).squeeze(1).detach().cpu().numpy()
|
|
if use_delta_encoding:
|
|
predicted_data = delta_decode(predicted_data)
|
|
|
|
# Ensure predicted_data is a flat list of floats
|
|
predicted_data = predicted_data[:min_length]
|
|
|
|
delta_data = [file_data[i].item() - predicted_data[i] for i in range(min_length)]
|
|
all_deltas.extend(delta_data)
|
|
|
|
if i == (epoch % len(data)):
|
|
visualize_prediction(file_data.cpu().numpy(), predicted_data, delta_data, sample_rate, epoch=epoch)
|
|
|
|
identical_percentage = (identical_count / len(data)) * 100
|
|
delta_plot_path = plot_delta_distribution(all_deltas, epoch)
|
|
wandb.log({"delta_distribution": wandb.Image(delta_plot_path)}, step=epoch)
|
|
|
|
return compression_ratios, identical_percentage
|
|
|
|
def train_model(model, train_data, test_data, epochs, batch_size, learning_rate, use_delta_encoding, encoder, eval_freq, save_path):
|
|
wandb.init(project="wav-compression")
|
|
criterion = nn.MSELoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
|
|
best_test_score = float('inf')
|
|
|
|
model.to(model.device)
|
|
|
|
max_length = max([len(seq) for seq in train_data])
|
|
print(f"Max sequence length: {max_length}")
|
|
|
|
for epoch in range(epochs):
|
|
total_loss = 0
|
|
random.shuffle(train_data)
|
|
for i in range(0, len(train_data) - batch_size, batch_size):
|
|
batch_data = [pad_sequence(np.array(train_data[j]), max_length) for j in range(i, i+batch_size)]
|
|
batch_data = np.array(batch_data)
|
|
inputs = torch.tensor(batch_data, dtype=torch.float32).unsqueeze(2).to(model.device)
|
|
targets = torch.tensor(batch_data, dtype=torch.float32).unsqueeze(2).to(model.device)
|
|
outputs = model(inputs)
|
|
loss = criterion(outputs, targets)
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
total_loss += loss.item()
|
|
|
|
wandb.log({"epoch": epoch, "loss": total_loss}, step=epoch)
|
|
print(f'Epoch {epoch+1}/{epochs}, Loss: {total_loss}')
|
|
|
|
if (epoch + 1) % eval_freq == 0:
|
|
train_compression_ratios, train_identical_percentage = evaluate_model(model, train_data, use_delta_encoding, encoder, epoch=epoch)
|
|
test_compression_ratios, test_identical_percentage = evaluate_model(model, test_data, use_delta_encoding, encoder, epoch=epoch)
|
|
|
|
wandb.log({
|
|
"train_compression_ratio_mean": np.mean(train_compression_ratios),
|
|
"train_compression_ratio_std": np.std(train_compression_ratios),
|
|
"train_compression_ratio_min": np.min(train_compression_ratios),
|
|
"train_compression_ratio_max": np.max(train_compression_ratios),
|
|
"test_compression_ratio_mean": np.mean(test_compression_ratios),
|
|
"test_compression_ratio_std": np.std(test_compression_ratios),
|
|
"test_compression_ratio_min": np.min(test_compression_ratios),
|
|
"test_compression_ratio_max": np.max(test_compression_ratios),
|
|
"train_identical_percentage": train_identical_percentage,
|
|
"test_identical_percentage": test_identical_percentage,
|
|
}, step=epoch)
|
|
|
|
print(f'Epoch {epoch+1}/{epochs}, Train Compression Ratio: Mean={np.mean(train_compression_ratios)}, Std={np.std(train_compression_ratios)}, Min={np.min(train_compression_ratios)}, Max={np.max(train_compression_ratios)}, Identical={train_identical_percentage}%')
|
|
print(f'Epoch {epoch+1}/{epochs}, Test Compression Ratio: Mean={np.mean(test_compression_ratios)}, Std={np.std(test_compression_ratios)}, Min={np.min(test_compression_ratios)}, Max={np.max(test_compression_ratios)}, Identical={test_identical_percentage}%')
|
|
|
|
test_score = np.mean(test_compression_ratios)
|
|
if test_score < best_test_score:
|
|
best_test_score = test_score
|
|
model_path = os.path.join(save_path, f"best_model_epoch_{epoch+1}.pt")
|
|
encoder_path = os.path.join(save_path, f"best_encoder_epoch_{epoch+1}.pkl")
|
|
torch.save(model.state_dict(), model_path)
|
|
with open(encoder_path, 'wb') as f:
|
|
pickle.dump(encoder, f)
|
|
print(f'New highscore on test data! Model and encoder saved to {model_path} and {encoder_path}')
|