PyTorch로 살펴보는 시간 마스킹 심층 분석

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import wandb
import numpy as np
import torchvision.transforms as transforms
import matplotlib.pyplot as plt


wandb.init(project="time_masking_experiments", name="random_masking")

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=5):
        super(SimpleCNN, self).__init__()
        
        # Convert the 2D spectrogram data to a 3D tensor: [batch_size, 1, freq_bins, time_frames]
        self.unsqueeze = lambda x: x.unsqueeze(1)
        
        self.conv1 = nn.Conv2d(1, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(2)
        self.flatten = nn.Flatten()
        
        # Calculate the output shape of the conv layers to adjust the FC layers
        self.sample_shape_after_convs = self._get_shape_after_convs(torch.zeros((1, 65, 77)).float())

        
        self.fc1 = nn.Linear(self.sample_shape_after_convs, 128)
        self.fc2 = nn.Linear(128, num_classes)

    def _get_shape_after_convs(self, x):
        x = self.unsqueeze(x)
        x = self.conv1(x)
        x = self.maxpool(x)
        x = self.conv2(x)
        x = self.maxpool(x)
        return np.prod(x.shape[1:])  # Multiply all dimensions except batch_size

    def forward(self, x):
        x = self.unsqueeze(x)  # Add channel dimension
        
        x = self.relu(self.conv1(x))
        x = self.maxpool(x)
        
        x = self.relu(self.conv2(x))
        x = self.maxpool(x)
        
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        
        return x

def random_masking(spec, T=40):
    # Clone the spectrogram to avoid in-place modifications
    masked_spectrogram = spec.clone()

    # Get the length of the time dimension
    time_length = masked_spectrogram.shape[2]

    # Determine the width of the mask
    mask_width = torch.randint(0, T + 1, (1,)).item()

    # Determine the starting point of the mask
    mask_start = torch.randint(0, time_length - mask_width + 1, (1,)).item()

    # Apply the mask
    masked_spectrogram[:, :, mask_start:mask_start + mask_width] = 0

    return masked_spectrogram

def frequency_time_masking(spectrogram, max_freq_mask_width=15, max_time_mask_width=40):
    # Clone the spectrogram to avoid in-place modifications
    masked_spectrogram = spectrogram.clone()

    # Apply frequency masking
    num_frequency_bins = masked_spectrogram.shape[1]
    freq_mask_width = torch.randint(0, max_freq_mask_width + 1, (1,)).item()
    freq_mask_start = torch.randint(0, num_frequency_bins - freq_mask_width + 1, (1,)).item()
    masked_spectrogram[:, freq_mask_start:freq_mask_start + freq_mask_width, :] = 0

    # Apply time masking
    num_time_frames = masked_spectrogram.shape[2]
    time_mask_width = torch.randint(0, max_time_mask_width + 1, (1,)).item()
    time_mask_start = torch.randint(0, num_time_frames - time_mask_width + 1, (1,)).item()
    masked_spectrogram[:, :, time_mask_start:time_mask_start + time_mask_width] = 0

    return masked_spectrogram

def generate_sine_wave(freq, time, sample_rate=5000):
    """Generate a sine wave of a given frequency."""
    t = np.linspace(0, time, int(time * sample_rate), endpoint=False)
    return np.sin(2 * np.pi * freq * t)

def generate_spectrogram(signal, sample_rate=5000, window_size=128, step_size=64):
    """Generate a spectrogram from a signal."""
    window = np.hanning(window_size)
    return plt.specgram(signal, NFFT=window_size, Fs=sample_rate, noverlap=window_size - step_size, window=window, mode='magnitude')[0]

class StructuredSpectrogramDataset(data.Dataset):
    def __init__(self, num_samples=100): # Note the change here
        self.num_samples = num_samples
        self.data = []
        self.labels = []

        for i in range(num_samples):
            freq = np.random.choice(np.arange(10, 2000, 10))  # Frequency up to 2000Hz

            # Remap the frequencies to 5 classes
            label = freq // 200  # This will automatically map to 0-4 for freq from 10-1000, and then 5-9 for freq from 1000-2000
            label %= 5  # This will make sure label is between 0-4

            sine_wave = generate_sine_wave(freq, 1)  # 1-second sine wave
            spectrogram = generate_spectrogram(sine_wave)
            self.data.append(spectrogram)
            self.labels.append(label)

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        return torch.tensor(self.data[idx], dtype=torch.float32).float(), self.labels[idx]

combined_dataset = StructuredSpectrogramDataset()

# Randomly shuffle and split
train_size = int(0.8 * len(combined_dataset))
eval_size = len(combined_dataset) - train_size

train_dataset, eval_dataset = torch.utils.data.random_split(combined_dataset, [train_size, eval_size])

# Create the dataloaders
train_loader = data.DataLoader(train_dataset, batch_size=16, shuffle=True)
eval_loader = data.DataLoader(eval_dataset, batch_size=16, shuffle=False)

def log_first_n_samples_to_wandb(n, data_loader):
    spectrogram_images = []
    for i, (data, target) in enumerate(data_loader):
        if i >= n:
            break
        # Convert tensor to numpy for visualization
        image = np.squeeze(data[0].numpy())
        spectrogram_images.append(wandb.Image(image, caption=f"Label: {target[0]}"))

    wandb.log({"First 5 Spectrograms": spectrogram_images})

log_first_n_samples_to_wandb(5, train_loader)

def train_one_epoch(model, data_loader, optimizer, criterion):
    model.train()
    total_loss = 0.0

    for batch_idx, (data, target) in enumerate(data_loader):
#         print("Shape of data:", data.shape)  # Add this line to print shape
        # Apply random time masking
        masked_data = random_masking(data)
        masked_data = masked_data.float()
        target = target.long() # CrossEntropyLoss expects targets to be long type



        # Log masked spectrograms to W&B (logging first sample in each batch)
        if batch_idx % 10 == 0:  # Log every 10 batches
            original_img = np.squeeze(data[0].numpy())
            masked_img = np.squeeze(masked_data[0].numpy())
            wandb.log({
                "Original Spectrogram": [wandb.Image(original_img, caption="Original")],
                "Masked Spectrogram": [wandb.Image(masked_img, caption="Masked")]
            })

        optimizer.zero_grad()
        outputs = model(masked_data)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    avg_loss = total_loss / len(data_loader)
    return avg_loss

def evaluate(model, eval_loader, criterion):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(eval_loader):
            outputs = model(data)
            loss = criterion(outputs, target)
            total_loss += loss.item()

            _, predicted = outputs.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()

    avg_loss = total_loss / len(eval_loader)
    accuracy = 100. * correct / total
    return avg_loss, accuracy

model = SimpleCNN(num_classes=5)
model = model.float()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

num_epochs = 10
for epoch in range(num_epochs):
    train_loss = train_one_epoch(model, train_loader, optimizer, criterion)
    eval_loss, eval_accuracy = evaluate(model, eval_loader, criterion)
    print(f"Epoch [{epoch+1}/{num_epochs}], Train Loss: {train_loss:.4f}, Eval Loss: {eval_loss:.4f}, Eval Accuracy: {eval_accuracy:.2f}%")
    
    # Log both Training and Evaluation Metrics
    wandb.log({
        "Train Loss": train_loss,
        "Eval Loss": eval_loss,
        "Eval Accuracy": eval_accuracy
    })

wandb.finish()