training.py

import torch
import torch.optim as optim
import customAudioDataset as data
import os
import torch.backends.cudnn as cudnn

from model import EncodecModel 
from msstftd import MultiScaleSTFTDiscriminator
from audio_to_mel import Audio2Mel

EPSILON = 1e-8
BATCH_SIZE = 5 #5#55
TENSOR_CUT = 50000 #10000
MAX_EPOCH = 10000 # Just set this to a very big number and manually stop it
SAVE_FOLDER = f'saves/new7/'
SAVE_LOCATION = f'{SAVE_FOLDER}batch{BATCH_SIZE}_cut{TENSOR_CUT}_' # appends epoch{epoch}.pth

if not os.path.exists(SAVE_FOLDER):
   os.makedirs(SAVE_FOLDER)

def total_loss(fmap_real, logits_fake, fmap_fake, wav1, wav2, sample_rate=24000):
    relu = torch.nn.ReLU()
    l1Loss = torch.nn.L1Loss(reduction='mean')
    l2Loss = torch.nn.MSELoss(reduction='mean')
    loss = torch.tensor([0.0], device='cuda', requires_grad=True)
    factor = 100 / (len(fmap_real) * len(fmap_real[0]))

    for tt1 in range(len(fmap_real)):
        loss = loss + (torch.mean(relu(1 - logits_fake[tt1])) / len(logits_fake))
        for tt2 in range(len(fmap_real[tt1])):
            loss = loss + (l1Loss(fmap_real[tt1][tt2].detach(), fmap_fake[tt1][tt2]) * factor)
    loss = loss * (2/3)

    for i in range(5, 11):
        fft = Audio2Mel(win_length=2 ** i, hop_length=2 ** i // 4, n_mel_channels=64, sampling_rate=sample_rate)
        loss = loss + l1Loss(fft(wav1), fft(wav2)) + l2Loss(fft(wav1), fft(wav2))
    loss = (loss / 6) + l1Loss(wav1, wav2)
    return loss

def disc_loss(logits_real, logits_fake):
    cx = torch.nn.ReLU()
    lossd = torch.tensor([0.0], device='cuda', requires_grad=True)
    for tt1 in range(len(logits_real)):
        lossd = lossd + torch.mean(cx(1-logits_real[tt1])) + torch.mean(cx(1+logits_fake[tt1]))
    lossd = lossd / len(logits_real)
    return lossd

def pad_sequence(batch):
    # Make all tensor in a batch the same length by padding with zeros
    batch = [item.permute(1, 0) for item in batch]
    batch = torch.nn.utils.rnn.pad_sequence(batch, batch_first=True, padding_value=0.)
    batch = batch.permute(0, 2, 1)
    return batch


def collate_fn(batch):
    B = len(batch)
    
    # rpsody: [1, 329, 1024]
    # timbre: [1, 512]
    # target: [1, 329, 128]

    max_length = max([item[1].shape[-2] for item in batch])
    # tar_max_length = max([item[3].shape[-2] for item in batch])

    pro_nfeats = batch[0][1].shape[-1]
    tim_nfeats = 512 # batch[0][2].shape[-1]
    tar_nfeats = batch[0][3].shape[-1]

    # print(type(B), type(pro_max_length), type(pro_nfeats))
    pro = torch.zeros((B, max_length, pro_nfeats), dtype=torch.float32)
    tar = torch.zeros((B, max_length, tar_nfeats), dtype=torch.float32)
    tim = torch.zeros((B, max_length, tim_nfeats), dtype=torch.float32)

    lengths = []

    for i, item in enumerate(batch):
        pro_, tim_, tar_ = item[1], item[2], item[3]
        
        lengths.append(pro_.shape[-2])
        # tar_lengths.append(tar_.shape[-2])

        pro[i,:pro_.shape[-2],:] = pro_
        tar[i,:tar_.shape[-2],:] = tar_
        tim[i,:tar_.shape[-2],:] = torch.rand(1, 512).unsqueeze(1).expand(-1,tar_.shape[-2],-1) # tim
    
    # tim = tim.expand(-1, )
    lengths = torch.LongTensor(lengths)
    # tar_lengths = torch.LongTensor(tar_lengths)

    print(pro.shape, tim.shape, tar.shape, lengths)
    return pro, tim, tar, lengths


def training(max_epoch = 5, log_interval = 20, fixed_length = 0, tensor_cut=100000, batch_size=8):
    # csv_path = 'datasets/e-gmd-v1.0.0/fileTRAIN.csv'
    # data_path = 'datasets/e-gmd-v1.0.0'
    # params.data_path 
    data_path = 'train.txt' # <uid>|<prosody_path>|<timbre_path>|<target_path>
    # timbre_path = 'datasets/timbre_features.txt' # <uid>|<path>
    # target_path = 'datasets/target_features.txt' # <uid>|<path>
    if fixed_length > 0:
        trainset = data.CustomAudioDataset(data_path, tensor_cut=tensor_cut, fixed_length=fixed_length)
    else:
        trainset = data.CustomAudioDataset(data_path, tensor_cut=tensor_cut)
    

    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn,)

    cudnn.benchmark = True

    target_bandwidths = [1.5, 3., 6, 12., 24.]
    sample_rate = 24_000
    channels = 1
    model = EncodecModel._get_model(
                target_bandwidths, sample_rate, channels,
                causal=False, model_norm='time_group_norm', audio_normalize=True,
                segment=1., name='disentangle_encodec_24khz')
    model.train()
    model.train_quantization = False
    model.cuda()
    
    # disc = MultiScaleSTFTDiscriminator(filters=32)
    # disc.train()
    # disc.cuda()


    lr = 0.01
    # optimizer = optim.SGD([{'params': model.parameters(), 'lr': lr}], momentum=0.9)
    # optimizer_disc = optim.SGD([{'params': disc.parameters(), 'lr': lr*10}], momentum=0.9)
    optimizer = optim.AdamW([{'params': model.parameters(), 'lr': lr}], betas=(0.8, 0.99))
    # optimizer_disc = optim.AdamW([{'params': disc.parameters(), 'lr': lr}], betas=(0.8, 0.99))

    def train(epoch):
        last_loss = 0
        # train_d = False
        print('----------------------------------------Epoch: {}----------------------------------------'.format(epoch))
        for batch_idx, batch in enumerate(trainloader):

            pro, tim, tar, lengths = batch
            # torch.Size([5, 754, 1024]) torch.Size([5]) torch.Size([5, 512]) torch.Size([5, 754, 128]) torch.Size([5]
            # torch.Size([5, 1376, 1024]) tensor([ 259,   98, 1376,  598, 1250]) torch.Size([5, 512]) torch.Size([5, 1376, 128]) tensor([ 259,   98, 1376,  598, 1250])
            pro = pro.cuda()
            lengths = lengths.cuda()
            tim = tim.cuda()
            # tim_lengths = tim_lengths.cuda()
            tar = tar.cuda()
            # tar_lengths = tar_lengths.cuda()


            optimizer.zero_grad()
            model.zero_grad()
            # optimizer_disc.zero_grad()
            # disc.zero_grad()
            print(tim.shape, pro.shape, lengths)
            # torch.Size([5, 643, 512]) torch.Size([5, 643, 1024]) tensor([102, 169, 643, 164, 319], device='cuda:0')
            output, loss_enc, _ = model(pro, tim, lengths)

            # logits_real, fmap_real = disc(input_wav)
            # if train_d:
            #     logits_fake, _ = disc(model(input_wav)[0].detach())
            #     loss = disc_loss(logits_real, logits_fake)
            #     if loss > last_loss/2:
            #         loss.backward()
            #         optimizer_disc.step()
            #     last_loss = 0

            # logits_fake, fmap_fake = disc(output)
            loss = total_loss(output, tim, tim_lengths)
            last_loss += loss.item()
            # loss_enc.backward(retain_graph=True)
            loss.backward()
            optimizer.step()

            if batch_idx % log_interval == 0:
                print(torch.cuda.mem_get_info())
                print(f"Train Epoch: {epoch} [{batch_idx * len(input_wav)}/{len(trainloader.dataset)} ({100. * batch_idx / len(trainloader):.0f}%)]")


    def adjust_learning_rate(optimizer, epoch):
        if epoch % 80 == 0:
            for param_group in optimizer.param_groups:
                param_group['lr'] = param_group['lr'] * 0.1


    for epoch in range(1, max_epoch):

        train(epoch)
        torch.save(model.state_dict(), f'{SAVE_LOCATION}epoch{epoch}.pth') #epoch{epoch}.pth
        # torch.save(disc.state_dict(), f'{SAVE_LOCATION}epoch{epoch}_disc.pth')

        adjust_learning_rate(optimizer, epoch)
        # adjust_learning_rate(optimizer_disc, epoch)

training(max_epoch=MAX_EPOCH, log_interval=100, fixed_length=0, batch_size=BATCH_SIZE, tensor_cut=TENSOR_CUT)