train.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-
from model import WaveRNN
import torch
from torch.autograd import Variable
import numpy as np
import time
from torch import optim
import torch.nn.functional as F
import sys
from hparam import Hyperparams as hp
from utils import split_signal
import math


def sine_wave(freq, length, sample_rate=hp.sr):
    return np.sin(np.arange(length) * 2 * math.pi * freq / sample_rate).astype(np.float32)


model = WaveRNN().cuda()
x = sine_wave(freq=500, length=hp.sr * 30)


def input_split():
    coarse_classes, fine_classes = split_signal(x)
    coarse_classes = np.reshape(coarse_classes, (1, -1))
    fine_classes = np.reshape(fine_classes, (1, -1))
    return coarse_classes, fine_classes


coarse_classes, fine_classes = input_split()


def train(model, optimizer, coarse_classes, fine_classes,num_steps, seq_len=960):
    start = time.time()
    running_loss = 0

    for step in range(num_steps):
        optimizer.zero_grad()
        loss = wavernn_loss(model, coarse_classes, fine_classes, seq_len)

        running_loss += (loss.data[0] / seq_len)
        loss.backward()
        optimizer.step()

        speed = (step + 1) / (time.time() - start)

        sys.stdout.write('\rStep: %i/%i --- NLL: %.2f --- Speed: %.3f batches/second ' %
                         (step + 1, num_steps, running_loss / (step + 1), speed))


def wavernn_loss(model, coarse_classes, fine_classes, seq_len):
    loss = 0
    hidden = model.init_hidden()
    rand_idx = np.random.randint(0, coarse_classes.shape[1] - seq_len - 1)
    for i in range(seq_len):
        j = rand_idx + i

        x_coarse = coarse_classes[:, j:j + 1]
        x_fine = fine_classes[:, j:j + 1]
        x_input = np.concatenate([x_coarse, x_fine], axis=1)
        x_input = x_input / 127.5 - 1.
        x_input = Variable(torch.FloatTensor(x_input)).cuda()

        y_coarse = coarse_classes[:, j + 1]
        y_fine = fine_classes[:, j + 1]
        y_coarse = Variable(torch.LongTensor(y_coarse)).cuda()
        y_fine = Variable(torch.LongTensor(y_fine)).cuda()

        current_coarse = y_coarse.float() / 127.5 - 1.
        current_coarse = current_coarse.unsqueeze(-1)

        out_coarse, out_fine, hidden = model(x_input, hidden, current_coarse)

        loss_coarse = F.cross_entropy(out_coarse, y_coarse)
        loss_fine = F.cross_entropy(out_fine, y_fine)
        loss += (loss_coarse + loss_fine)
    return loss


optimizer = optim.Adam(model.parameters(), lr=1e-3)
train(model, optimizer, coarse_classes, fine_classes, num_steps=500)