深入理解 PyTorch 的数据加载

在进行深度学习时，数据的加载和预处理是至关重要的步骤。PyTorch 提供了 torch.utils.data.Dataset 和 torch.utils.data.DataLoader 这两个强大的工具来简化这一过程。

1. torch.utils.data.Dataset

Dataset 是 PyTorch 中用于定义数据集的基类。通过继承 Dataset 类，您可以创建自己的数据集。实现 __len__ 和 __getitem__ 方法是创建自定义数据集的关键。

主要方法

• __len__：返回数据集的大小，表示数据集中样本的总数。
• __getitem__：根据索引返回数据集中的一个样本。在此方法中，您可以实现数据的读取和预处理逻辑。

注意事项

• 数据长度一致性：在处理音频或图像数据时，确保每个样本的长度一致是非常重要的。如果样本长度不一致，DataLoader 将无法将它们堆叠成批次，可能会导致运行时错误。
• 数据预处理：在 __getitem__ 方法中，您可以进行必要的预处理，例如归一化、数据增强等。

2. torch.utils.data.DataLoader

DataLoader 是 PyTorch 中用于加载数据的工具。它可以自动处理批次、打乱数据和多线程加载。

主要参数

• dataset：需要加载的数据集。
• batch_size：每个批次的样本数量。
• shuffle：在每个 epoch 之前是否打乱数据。
• num_workers：用于数据加载的子进程数量，可以加速数据加载。

注意事项

• 批次大小：选择合适的批次大小可以提高训练效率。过小的批次可能导致训练不稳定，而过大的批次可能导致内存溢出。
• 数据打乱：在每个 epoch 之前打乱数据可以帮助模型更好地泛化。

代码示例

以下是一个完整的示例代码，展示了如何使用 torch.utils.data.Dataset 和 torch.utils.data.DataLoader 来处理音频数据。

import os
import random
import torch
from torch.utils.data import Dataset, DataLoader
import torchaudio
import torch.nn.functional as F
import librosa
import numpy as npclass AudioDataset(Dataset):def __init__(self, mix_dir, noise_dir, voice_dir, target_length=1024, num_samples=5):self.mix_dir = mix_dirself.noise_dir = noise_dirself.voice_dir = voice_dirself.target_length = target_lengthself.num_samples = num_samples  # 每个音频文件提取的样本数量# 获取所有混合音频文件名（去掉 '_mix.wav' 后缀）self.file_names = [f.replace('_mix.wav', '') for f in os.listdir(mix_dir) if f.endswith('_mix.wav')]def __len__(self):return len(self.file_names)def __getitem__(self, idx):# 获取文件名file_name = self.file_names[idx]# 构建文件路径mix_path = os.path.join(self.mix_dir, f"{file_name}_mix.wav")noise_path = os.path.join(self.noise_dir, f"{file_name}_noise.wav")voice_path = os.path.join(self.voice_dir, f"{file_name}_voice.wav")# 读取音频文件mix, _ = torchaudio.load(mix_path)noise, _ = torchaudio.load(noise_path)voice, _ = torchaudio.load(voice_path)# 随机选择多个 1024 个数据点的片段mix_samples = self._get_random_samples(mix)noise_samples = self._get_random_samples(noise)voice_samples = self._get_random_samples(voice)# 计算 STFT 并返回频谱mix_spectrogram = [self._compute_stft(sample) for sample in mix_samples]noise_spectrogram = [self._compute_stft(sample) for sample in noise_samples]voice_spectrogram = [self._compute_stft(sample) for sample in voice_samples]return mix_spectrogram, noise_spectrogram, voice_spectrogramdef _get_random_samples(self, audio):samples = []for _ in range(self.num_samples):if audio.size(1) > self.target_length:start_idx = random.randint(0, audio.size(1) - self.target_length)sample = audio[:, start_idx:start_idx + self.target_length]else:# 如果音频长度小于目标长度，则填充padding = self.target_length - audio.size(1)sample = F.pad(audio, (0, padding), value=0)  # 在末尾填充零samples.append(sample)return samples@staticmethoddef _compute_stft(audio):# 将张量转换为 NumPy 数组audio_np = audio.numpy().flatten()  # 转换为一维数组# 计算 STFTstft_result = librosa.stft(audio_np, n_fft=512, hop_length=256)# 转换为幅度谱spectrogram = np.abs(stft_result)return spectrogram  # 返回幅度谱# 使用示例
if __name__ == "__main__":# 文件夹路径mix_directory = './ds/mixtures'noise_directory = './ds/noises'voice_directory = './ds/real_voices'# 创建数据集audio_dataset = AudioDataset(mix_directory, noise_directory, voice_directory, num_samples=5)# 创建数据加载器data_loader = DataLoader(audio_dataset, batch_size=16, shuffle=True)# 遍历数据for mix_spectrogram, noise_spectrogram, voice_spectrogram in data_loader:# 将样本在第一维拼接mix_tensor = torch.cat(mix_spectrogram, dim=0)  # 拼接 mix_samplesnoise_tensor = torch.cat(noise_spectrogram, dim=0)  # 拼接 noise_samplesvoice_tensor = torch.cat(voice_spectrogram, dim=0)  # 拼接 voice_samples# 处理您的频谱数据print(f'Mix spectrogram shape: {mix_tensor.shape}')  # 打印每个频谱的形状print(f'Noise spectrogram shape: {noise_tensor.shape}')  # 打印每个频谱的形状print(f'Voice spectrogram shape: {voice_tensor.shape}')  # 打印每个频谱的形状

Mix spectrogram shape: torch.Size([80, 257, 5])
Noise spectrogram shape: torch.Size([80, 257, 5])
Voice spectrogram shape: torch.Size([80, 257, 5])
Mix spectrogram shape: torch.Size([80, 257, 5])
Noise spectrogram shape: torch.Size([80, 257, 5])
Voice spectrogram shape: torch.Size([80, 257, 5])
Mix spectrogram shape: torch.Size([40, 257, 5])
Noise spectrogram shape: torch.Size([40, 257, 5])
Voice spectrogram shape: torch.Size([40, 257, 5])