import torch.nn as nn
import torch
import numpy as np
import torch.nn.functional as F
from hparams import hparams
from dataset import WavNet_Dataset
from torch.utils.data import Dataset,DataLoader
from model import WavNet
import logging
import os
def adjust_lr_rate(optimizer,lr,lr_decay):
 lr_new = max(0.00005, lr - lr_decay)
 
 for param_groups in optimizer.param_groups:
 param_groups['lr'] = lr_new
 return lr_new,optimizer
 

if __name__ == "__main__":
 
 # 定义log文件
 file_log = "WaveNet.log"
 logging.basicConfig(
 level=logging.INFO,
 format='%(asctime)s - %(levelname)s - %(message)s',
 handlers=[
 logging.FileHandler(file_log),
 logging.StreamHandler()
 ]
 )
 logger = logging.getLogger() 

 # 定义device
 device = torch.device("cuda:0")
 
 # 获取模型参数
 para = hparams()
 
 # 模型实例化
 m_model = WavNet(para)
 m_model = m_model.to(device)
 receptive_field = m_model.receptive_field
 # 定义优化器
 m_optimizer = torch.optim.Adam(m_model.parameters(), para.lr, [0.5, 0.999])
 lr = para.lr
 # # 模型加载
 # model_name = 'save/10/model.pick'
 # m_model_all = torch.load(model_name)
 # m_model.load_state_dict(m_model_all['model'])
 # m_optimizer.load_state_dict(m_model_all['opt'])
 
 
 
 
 
 
 
 
 # 损失函数
 CELoss = nn.CrossEntropyLoss()
 
 # 定义数据集 
 m_Dataset= WavNet_Dataset(para,receptive_field,para.output_length)
 print(len(m_Dataset))
 m_DataLoader = DataLoader(m_Dataset,batch_size = para.batch_size,shuffle = True, num_workers = 8)
 
 
 # 开始训练 
 for epoch in range(para.n_epoch):
 # 调整lr
 if epoch>para.start_decay and (epoch-para.start_decay)%(para.lr_update_epoch)==0:
 lr, m_optimizer= adjust_lr_rate(m_optimizer,lr,para.decay_lr)
 
 for i, sample_batch in enumerate(m_DataLoader):
 loss = []
 train_data = sample_batch[0]
 train_target = sample_batch[1]
 train_data = train_data.to(device)
 train_target = train_target.to(device) # [B,out_length]
 
 outputs = m_model(train_data) # [B,C,out_length]
 outputs = outputs.transpose(1,2).contiguous().view(-1,para.classes) #[B*out_length,C]
 train_target = train_target.view(-1) # [B*out_length]
 
 # 计算损失函数
 m_loss = CELoss(outputs,train_target)
 
 m_optimizer.zero_grad()
 m_loss.backward()
 m_optimizer.step()
 
 loss.append(m_loss.cpu().detach().numpy() )
 
 # log 输出
 logger.info("epoch %8d step %8d loss= %f"%(epoch,i,m_loss))
 
 # 保存模型
 path_save = os.path.join(para.path_save,str(epoch))
 os.makedirs(path_save,exist_ok=True)
 
 torch.save({'model':m_model.state_dict(),
 'opt':m_optimizer.state_dict()},
 os.path.join(path_save,'model.pick'))
 
 logger.info("epoch %8d loss_mean= %f"%(epoch,np.mean(loss)))