HiFi-GAN

bin:
    download.py  # Скрипт для загрузки данных
notebooks:
    checks.ipynb  # Проверка функционала
    datasphere.ipynb  # Запуск скриптов в ДатаСфере
resources:
    chkpoints  # Полные чекпоинты для перезапуска экспериментов
    ids  # ID запусков W&B
    models  # Чекпоинты генератора, отобранные по метрике на валидации
    predicted  # Результаты инференса
    test_audio  # Фудиофайлы для инференса
src:
    config  # Файлы конфигурации, содержащие все параметры экспериментов
    dataset  # Загрузчики данных
    loops  # Циклы обучения, валидации и т.д.
    loss  # Функции потерь
    models  # Архитектуры моделей
    transforms  # Преобразования аудио (мел спектрограммы и т.д.)
    utils  # Утилиты (сохранение и загрузка чекпоинтов, визуализация и т.д.)
train.py  # Скрипт для обучения
inference.py  # Скрипт для инференса

def __getitem__(self, i: int) -> Tensor:
    audio, sample_rate = torchaudio.load(self.ids[i])
    assert sample_rate == self.sample_rate
    audio = audio / audio.max(dim=1, keepdim=True)[0] * self.wav_scale

    if self.segment_size is not None and not self.validation:
        if audio.shape[1] >= self.segment_size:
            audio_start_bound = audio.shape[1] - self.segment_size
            audio_start = np.random.randint(audio_start_bound)
            audio = audio[:, audio_start:audio_start + self.segment_size]
        else:
            # I'm almost 100% sure this is never used
            audio = F.pad(audio, (0, self.segment_size - audio.shape[1]), "constant")

    if self.augmentations is not None:
        audio = self.augmentations(audio)

upsample_rates = [8, 8, 2, 2]
upsample_kernel_sizes = [16, 16, 4, 4]
upsample_initial_channel = 512
resblock_kernel_sizes = [3, 7, 11]
resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]

# Parameters from the official implementation
mpd_periods = [2, 3, 5, 7, 11]
mpd_channels = [1, 32, 128, 512, 1024, 1024]
mpd_kernel_size = 5
mpd_stride = 3
mpd_use_spectral_norm = False

class MultiPeriodDiscriminator(nn.Module):
    def __init__(self, config):
        super().__init__()
        discriminators = [PeriodDiscriminator(config, period)
                          for period in config.mpd_periods]
        self.discriminators = nn.ModuleList(discriminators)

    def forward(self, real: Tensor, fake: Tensor) -> Dict[str, List]:
        # ... do something ...
        return {"mpd_real_outs": real_outs,
                "mpd_fake_outs": fake_outs,
                "mpd_real_features": real_features,
                "mpd_fake_features": fake_features}

# Parameters from the official implementation
msd_channels = [1, 128, 128, 256, 512, 1024, 1024, 1024]
msd_kernel_sizes = [15, 41, 41, 41, 41, 41, 5]
msd_strides = [1, 2, 2, 4, 4, 1, 1]
msd_groups = [1, 4, 16, 16, 16, 16, 1]

class Discriminator(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.mpd = MultiPeriodDiscriminator(config)
        self.msd = MultiScaleDiscriminator(config)

    def forward(self, real: Tensor, fake: Tensor) -> Dict[str, List]:
        mpd_out = self.mpd(real, fake)
        msd_out = self.msd(real, fake)
        return {**mpd_out, **msd_out}

def _apply_conv_hook(function, module: nn.Module):
    name = module.__class__.__name__
    if "Conv" in name:
        function(module)

def conv_hook(function):
    return partial(_apply_conv_hook, function)

class DiscriminatorLoss(nn.Module):
    def forward(self, real_outs: List[Tensor], fake_outs: List[Tensor]) -> Tensor:
        loss = 0
        for real, fake in zip(real_outs, fake_outs):
            # Classify real as 1 (real), fake as 0 (fake)
            loss = loss + ((1 - real) ** 2).mean() + (fake ** 2).mean()
        return loss

class FeatureLoss(nn.Module):
    def forward(self, real_features: List[Tensor], fake_features: List[Tensor]) -> Tensor:
        loss = 0
        for real_feature_map, fake_feature_map in zip(real_features, fake_features):
            for real_feature, fake_feature in zip(real_feature_map, fake_feature_map):
                loss = loss + self.l1_loss(fake_feature, real_feature)
        return loss

class GeneratorLoss(nn.Module):
    def forward(self, fake_outs: List[Tensor]) -> Tensor:
        loss = 0
        for fake in fake_outs:
            loss = loss + ((fake - 1) ** 2).mean()  # Classify fake as 1 (real)
        return loss

# Get generator predictions
pred_audio = generator(input_mel)
pred_mel = mel_spectrogram_loss(pred_audio.squeeze(1))

# Pad targets to match predictions from generator
target_audio = pad_to_length(target_audio, pred_audio.shape[-1])
target_mel = pad_to_length(target_mel, pred_mel.shape[-1])

# Discriminator step
dis_optimizer.zero_grad()

# Discriminator update
discriminator.requires_grad(True)

dis_out = discriminator(target_audio, pred_audio.detach())
dis_losses = compute_losses(dis_criterion, dis_out)
dis_loss = compute_total_loss(dis_losses)
dis_loss.backward()
dis_optimizer.step()

# Generator step
gen_optimizer.zero_grad()

# Turn off discriminator gradients
discriminator.requires_grad(False)

# Get discriminator predictions
dis_out = discriminator(target_audio, pred_audio)

# Generator update
super_losses = compute_losses(super_criterion, target_mel, pred_mel)
gen_losses = compute_losses(gen_criterion, dis_out)
gen_loss = compute_total_loss({**super_losses, **gen_losses})
gen_loss.backward()
gen_optimizer.step()