[DOCS FIX]

training/gan/gan.py

@@ -0,0 +1,167 @@
+import logging
+import os
+import shutil
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchaudio
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+
+logging.basicConfig(
+ level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+class Generator(nn.Module):
+ def __init__(self, latent_dim: int, output_dim: int):
+ super().__init__()
+ self.model = nn.Sequential(
+ nn.Linear(latent_dim, 256),
+ nn.LeakyReLU(0.2),
+ nn.Linear(256, 512),
+ nn.LeakyReLU(0.2),
+ nn.Linear(512, 1024),
+ nn.LeakyReLU(0.2),
+ nn.Linear(1024, output_dim),
+ nn.Tanh(),
+ )
+
+ def forward(self, z: torch.Tensor) -> torch.Tensor:
+ return self.model(z)
+
+
+class Discriminator(nn.Module):
+ def __init__(self, input_dim: int):
+ super().__init__()
+ self.model = nn.Sequential(
+ nn.Linear(input_dim, 1024),
+ nn.LeakyReLU(0.2),
+ nn.Linear(1024, 512),
+ nn.LeakyReLU(0.2),
+ nn.Linear(512, 256),
+ nn.LeakyReLU(0.2),
+ nn.Linear(256, 1),
+ nn.Sigmoid(),
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.model(x)
+
+
+def train_gan(
+ generator: nn.Module,
+ discriminator: nn.Module,
+ dataloader: DataLoader,
+ num_epochs: int,
+ latent_dim: int,
+ device: torch.device,
+) -> Tuple[nn.Module, nn.Module]:
+ criterion = nn.BCELoss()
+ g_optimizer = optim.Adam(
+ generator.parameters(), lr=0.0002, betas=(0.5, 0.999)
+ )
+ d_optimizer = optim.Adam(
+ discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999)
+ )
+
+ for epoch in range(num_epochs):
+ for i, batch in enumerate(dataloader):
+ real_samples = batch["audio"].to(device)
+ batch_size = real_samples.size(0)
+
+ d_optimizer.zero_grad()
+ real_labels = torch.ones(batch_size, 1).to(device)
+ fake_labels = torch.zeros(batch_size, 1).to(device)
+
+ d_real_output = discriminator(real_samples)
+ d_real_loss = criterion(d_real_output, real_labels)
+
+ z = torch.randn(batch_size, latent_dim).to(device)
+ fake_samples = generator(z)
+ d_fake_output = discriminator(fake_samples.detach())
+ d_fake_loss = criterion(d_fake_output, fake_labels)
+
+ d_loss = d_real_loss + d_fake_loss
+ d_loss.backward()
+ d_optimizer.step()
+
+ g_optimizer.zero_grad()
+ z = torch.randn(batch_size, latent_dim).to(device)
+ fake_samples = generator(z)
+ d_output = discriminator(fake_samples)
+ g_loss = criterion(d_output, real_labels)
+ g_loss.backward()
+ g_optimizer.step()
+
+ logger.info(
+ f"Epoch [{epoch+1}/{num_epochs}] D_loss: {d_loss.item():.4f} G_loss: {g_loss.item():.4f}"
+ )
+
+ return generator, discriminator
+
+
+def main():
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ logger.info(f"Using device: {device}")
+
+ # Clear the dataset cache
+ cache_dir = os.path.expanduser("~/.cache/huggingface/datasets")
+ if os.path.exists(cache_dir):
+ shutil.rmtree(cache_dir)
+ logger.info("Cleared dataset cache")
+
+ # Try to load the dataset with retries
+ max_retries = 3
+ for attempt in range(max_retries):
+ try:
+ dataset = load_dataset(
+ "mozilla-foundation/common_voice_11_0",
+ "en",
+ split="train[:1000]",
+ trust_remote_code=True,
+ )
+ logger.info(f"Dataset loaded: {len(dataset)} samples")
+ break
+ except Exception as e:
+ if attempt < max_retries - 1:
+ logger.warning(
+ f"Dataset loading failed (attempt {attempt + 1}/{max_retries}). Retrying..."
+ )
+ else:
+ logger.error("Failed to load dataset after multiple attempts.")
+ raise e
+
+ def preprocess_audio(example):
+ audio = example["audio"]["array"]
+ resampled_audio = torchaudio.transforms.Resample(
+ example["audio"]["sampling_rate"], 16000
+ )(torch.tensor(audio))
+ return {"audio": resampled_audio.flatten()}
+
+ dataset = dataset.map(preprocess_audio, remove_columns=dataset.column_names)
+ dataset.set_format(type="torch", columns=["audio"])
+
+ dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
+
+ latent_dim = 100
+ output_dim = 16000
+
+ generator = Generator(latent_dim, output_dim).to(device)
+ discriminator = Discriminator(output_dim).to(device)
+
+ num_epochs = 50
+ generator, discriminator = train_gan(
+ generator, discriminator, dataloader, num_epochs, latent_dim, device
+ )
+
+ torch.save(generator.state_dict(), "generator.pth")
+ torch.save(discriminator.state_dict(), "discriminator.pth")
+ logger.info("Models saved successfully")
+
+
+if __name__ == "__main__":
+ main()