<reponame>ifding/dynamic-analysis-firmware
"""
Main model of WaveNet
Calculate loss and optimizing
"""
import os

import torch
import torch.optim
import torch.nn.functional as F
from torch.autograd import Variable

from warpctc_pytorch import CTCLoss

from model.networks import WaveNet as WaveNetModule


class WaveNet:
    def __init__(self, layer_size, stack_size, in_channels, res_channels, lr=0.002):

        self.net = WaveNetModule(layer_size, stack_size, in_channels, res_channels)

        self.in_channels = in_channels
        self.receptive_fields = self.net.receptive_fields

        self.lr = lr
        self.loss = self._loss()
        self.optimizer = self._optimizer()
        self.criterion = CTCLoss()

        self._prepare_for_gpu()

    @staticmethod
    def _loss():
        loss = torch.nn.CrossEntropyLoss()

        if torch.cuda.is_available():
            loss = loss.cuda()

        return loss

    def _optimizer(self):
        return torch.optim.Adam(self.net.parameters(), lr=self.lr)

    def _prepare_for_gpu(self):
        #if torch.cuda.device_count() > 1:
        #    print("{0} GPUs are detected.".format(torch.cuda.device_count()))
        #    self.net = torch.nn.DataParallel(self.net)

        if torch.cuda.is_available():
            self.net.cuda()

    def train(self, inputs, targets, input_percentages, target_sizes):
        """
        Train 1 time
        :param inputs: Tensor[batch, timestep, channels]
        :param targets: Torch tensor [batch, timestep, channels]
        :return: float loss
        """
        
        inputs = Variable(inputs, requires_grad=False)
        target_sizes = Variable(target_sizes, requires_grad=False)
        targets = Variable(targets, requires_grad=False)        
        
        if torch.cuda.is_available():
            inputs = inputs.cuda()
        
        inputs = torch.squeeze(inputs)  
        
        outputs = self.net(inputs)   
        
        outputs = outputs.transpose(0, 1)  # TxNxH

        seq_length = outputs.size(0)
        sizes = Variable(input_percentages.mul_(int(seq_length)).int(), requires_grad=False) 

        loss = self.criterion(outputs, targets, sizes, target_sizes)   
        
        loss = loss / inputs.size(0)  # average the loss by minibatch
            
        loss_sum = loss.data.sum()
        inf = float("inf")
        if loss_sum == inf or loss_sum == -inf:
            print("WARNING: received an inf loss, setting loss value to 0")
            loss_value = 0
        else:
            loss_value = loss.item()
           
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        #if torch.cuda.is_available():
        #    torch.cuda.synchronize()        

        return loss_value
    
    def _train(self, inputs, targets):
        """
        Train 1 time
        :param inputs: Tensor[batch, timestep, channels]
        :param targets: Torch tensor [batch, timestep, channels]
        :return: float loss
        """
        outputs = self.net(inputs)

        loss = self.loss(outputs.view(-1, self.in_channels),
                         targets.long().view(-1))

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return loss.data[0]
    
    
    def generate(self, inputs):
        """
        Generate 1 time
        :param inputs: Tensor[batch, timestep, channels]
        :return: Tensor[batch, timestep, channels]
        """
        
        if torch.cuda.is_available():
            inputs = inputs.cuda()
            
        inputs = torch.squeeze(inputs)
        
        outputs = self.net(inputs)

        return outputs

    @staticmethod
    def get_model_path(model_dir, step=0):
        basename = 'wavenet'

        if step:
            return os.path.join(model_dir, '{0}_{1}.pkl'.format(basename, step))
        else:
            return os.path.join(model_dir, '{0}.pkl'.format(basename))

    def load(self, model_dir, step=0):
        """
        Load pre-trained model
        :param model_dir:
        :param step:
        :return:
        """
        print("Loading model from {0}".format(model_dir))

        model_path = self.get_model_path(model_dir, step)

        self.net.load_state_dict(torch.load(model_path))

    def save(self, model_dir, step=0):
        print("Saving model into {0}".format(model_dir))

        model_path = self.get_model_path(model_dir, step)

        torch.save(self.net.state_dict(), model_path)

