train_rpn_kws.py

#!/usr/bin/env python

# Copyrigh 2018 houjingyong@gmail.com

# MIT Licence

from __future__ import print_function

import os, sys, argparse, datetime, shutil
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

from bbox_transform import get_out_utt_boxes
from config import cfg
from config import cfg_from_file
from streaming_special_torch_dataset import * 
from kaldi_io import *
from RNNs import GRU
from RPN import RPN
from RPN_KWS import RPN_KWS
from utils import AverageMeter, count_parameters
from loss import loss_frame_fn_ce, acc_frame

def get_args():
    """Get arguments from stdin."""
    parser = argparse.ArgumentParser(description='Pytorch acoustic model.')
    parser.add_argument('--encoder', type=str, default='gru',
                        help='encoder type {default: gru}')
    parser.add_argument('--num-anchor', type=int, default=10, metavar='HF',
                        help='Num anchors per frame {default: 10.0}')
    parser.add_argument('--lambda-factor', type=float, default=5.0, metavar='HF',
                        help='Balance factor between classification and regression loss (default: 5.0).')
    parser.add_argument('--input-dim', type=int, default=40, metavar='N',
                        help='Input feature dimension without context (default: 40).')
    parser.add_argument('--kernel-size', type=int, default=3, metavar='N',
                        help='Kernel size of Wavenet or CNN (default:3).')
    parser.add_argument('--hidden-dim', type=int, default=128, metavar='N',
                        help='Hidden dimension of feature extractor (default: 128).')
    parser.add_argument('--num-layers', type=int, default=2, metavar='N',
                        help='Numbers of hidden layers of feature extractor (default: 2).')
    parser.add_argument('--output-dim', type=int, default=2000, metavar='N',
                        help='Output dimension, number of classes (default: 2000).')
    parser.add_argument('--dropout', type=float, default=0.0001, metavar='DR',
                        help='dropout of feature extractor (default: 0.0001).')
    parser.add_argument('--left-context', type=int, default=5, metavar='N',
                        help='Left context length for splicing feature (default: 5).')
    parser.add_argument('--right-context', type=int, default=5, metavar='N',
                        help='Right context length for splicing feature (default: 5).')
    parser.add_argument('--max-epochs', type=int, default=20, metavar='N',
                        help='Maximum epochs to train (default: 20).')
    parser.add_argument('--min-epochs', type=int, default=0, metavar='N',
                        help='Minimum epochs to train (default: 0).')
    parser.add_argument('--batch-size', type=int, default=8, metavar='N',
                        help='Batch size for training (default: 8).')
    parser.add_argument('--learning-rate', type=float, default=0.001, metavar='LR',
                        help='Initial learning rate (default: 0.001).')
    parser.add_argument('--halving-factor', type=float, default=0.5, metavar='HF',
                        help='Half factor for learning rate (default: 0.5).')
    parser.add_argument('--start-halving-impr', type=float, default=0.01, metavar='S',
                        help='Improvement threshold to half the learning rate (default: 0.01).')
    parser.add_argument('--end-halving-impr', type=float, default=0.001, metavar='E',
                        help='Improvement threshold to stop half learning rate (default: 0.001).')
    parser.add_argument('--init-weight-decay', type=float, default=1e-5, metavar='E',
                                    help='Weight decay of L2 normalization (default: 1e-5).')
    parser.add_argument('--seed', type=int, default=1234, metavar='S',
                        help='Random seed (default: 1234).')
    parser.add_argument('--use-cuda', type=int, default=1, metavar='C',
                        help='Use cuda (1) or cpu(0).')
    parser.add_argument('--multi-gpu', type=int, default=0, metavar='G',
                        help='Use multi gpu (1) or not (0).')
    parser.add_argument('--train', type=int, default=1,
                        help='Executing mode, train (1) or test (0).')
    parser.add_argument('--train-scp', type=str, default='',
                        help='Training data file.')
    parser.add_argument('--dev-scp', type=str, default='',
                        help='Development data file.')
    parser.add_argument('--save-dir', type=str, default='',
                        help='Directory to output the model.')
    parser.add_argument('--load-model', type=str, default='',
                        help='Previous model to load.')
    parser.add_argument('--test', type=int, default=0,
                        help='Executing mode, 1 for test, 0 no test')
    parser.add_argument('--test-scp', type=str, default='',
                        help='Test data file.')
    parser.add_argument('--output-file', type=str, default='',
                        help='Test output file')
    parser.add_argument('--region-output-file', type=str, default='',
                        help='Region output file')
    parser.add_argument('--log-interval', type=int, default=1000, metavar='N',
                        help='How many batches to wait before logging training status.')
    parser.add_argument('--num-workers', type=int, default=1, metavar='N',
                        help='How many workers used to load data')
    parser.add_argument('--config-file', type=str, default='',
                        help='config file in yaml format')
    args = parser.parse_args()

    if args.config_file != '':
        cfg_from_file(args.config_file)
    
    return args

def get_new_target(device, target, num_p, num_n):
    new_target=[]
    for i in range(target.size(0)):
        if target[i][0] == 0:
            new_target += ([target[i][0]] * num_n)
        else:
            new_target += ([target[i][0]] * num_p)
    return torch.LongTensor(new_target).to(device)

def adjust_learning_rate(args, optimizer):
    """Half the learning rate when relative improvement is too low.
    Args:
        args: Arguments for training.
        optimizer: Optimizer for training.
    """
    args.learning_rate *= args.halving_factor
    for param_group in optimizer.param_groups:
        param_group['lr'] = args.learning_rate

def train(args, model, device, train_loader, optimizer, epoch):
    """Train one epoch."""
    tr_rpn_loss_bbox = AverageMeter()
    tr_rpn_loss_cls = AverageMeter()
    tr_loss = AverageMeter()
    tr_rpn_acc = AverageMeter()
    model.train()
    total_step = len(train_loader)
    balance_weight=args.lambda_factor
    for batch_idx, (utt_id, act_lens, data, target) in enumerate(train_loader):
        act_lens, data, target = act_lens.to(device), data.to(device), target.to(device)
        target = target.reshape(target.size(0), 1, target.size(1)).float()
        # Forward pass
        batch_size = data.shape[0]
        outputs = model(epoch, data, act_lens, target, 100)
        rois, rpn_cls_score, rpn_label, rpn_loss_cls, rpn_loss_bbox = outputs
        rpn_acc = acc_frame(rpn_cls_score, rpn_label)
        
        # Backward and optimize
        loss = rpn_loss_cls + balance_weight * rpn_loss_bbox
        optimizer.zero_grad()
        loss.backward()
        #name, param=list(model.named_parameters())[1]
        #print('Epoch:[{}/{}], param name:{},\n param:'.format(epoch+1, args.max_epochs, name, param))
        optimizer.step()

        tr_rpn_acc.update(rpn_acc, 1)
        tr_loss.update(loss, 1)
        tr_rpn_loss_cls.update(rpn_loss_cls, 1)
        tr_rpn_loss_bbox.update(rpn_loss_bbox, 1)

        if batch_idx % args.log_interval == 0:
            print('Epoch: [{}/{}], Step [{}/{}], Train Loss: {:.4f}, Train RPN Acc: {:.4f}%'
                  .format(epoch+1, args.max_epochs, batch_idx+1, total_step, tr_loss.cur,  tr_rpn_acc.cur))
            print('Epoch: [{}/{}], Step [{}/{}], Train RPN cls Loss: {:.4f}, Train RPN bbox Loss: {:.4f} '
                  .format(epoch+1, args.max_epochs, batch_idx+1, total_step, tr_rpn_loss_cls.cur, tr_rpn_loss_bbox.cur))

    print('Epoch: [{}/{}], Average Train Loss: {:.4f}, Average Train RPN cls Loss: {:.4f}, Average Train RPN bbox Loss: {:.4f}, AverageAverage Train RPN Acc: {:.4f}%'
         .format(epoch+1, args.max_epochs, tr_loss.avg, tr_rpn_loss_cls.avg, tr_rpn_loss_bbox.avg, tr_rpn_acc.avg))
    return float("{:.4f}".format(tr_loss.avg))


def validate(args, model, device, dev_loader, epoch):
    """Cross validate the model."""
    meter_rpn_loss_bbox = AverageMeter()
    balance_weight = args.lambda_factor
    meter_rpn_loss_cls = AverageMeter()
    meter_loss = AverageMeter()
    meter_rpn_acc = AverageMeter()
    balance_weight = args.lambda_factor
    with torch.no_grad():
        total_step = len(dev_loader)
        for batch_idx, (utt_id, act_lens, data, target) in enumerate(dev_loader):
            act_lens, data, target = act_lens.to(device), data.to(device), target.to(device)
            target = target.reshape(target.size(0), 1, target.size(1)).float()
            # Forward pass
            batch_size = data.shape[0]
            outputs = model(epoch, data, act_lens, target, 100)
            rois, rpn_cls_score, rpn_label, rpn_loss_cls, rpn_loss_bbox = outputs
            rpn_acc = acc_frame(rpn_cls_score, rpn_label)
            # Backward and optimize
            loss = rpn_loss_cls + balance_weight * rpn_loss_bbox 
            meter_rpn_acc.update(rpn_acc, 1)
            meter_loss.update(loss, 1)
            meter_rpn_loss_cls.update(rpn_loss_cls, 1)
            meter_rpn_loss_bbox.update(rpn_loss_bbox, 1)

            if batch_idx % args.log_interval == 0:
                print('Epoch: [{}/{}], Step [{}/{}], Val Loss: {:.4f}, Val RPN Acc: {:.4f}% '
                      .format(epoch+1, args.max_epochs, batch_idx+1, total_step, meter_loss.cur, meter_rpn_acc.cur))
                print('Epoch: [{}/{}], Step [{}/{}], Val RPN cls Loss: {:.4f}, Val RPN bbox Loss: {:.4f} '
                      .format(epoch+1, args.max_epochs, batch_idx+1, total_step, meter_rpn_loss_cls.cur, meter_rpn_loss_bbox.cur))

        print('Epoch: [{}/{}], Average Val Loss: {:.4f}, Average Val RPN cls Loss: {:.4f}, Average Val RPN bbox Loss: {:.4f}, Average Val RPN Acc: {:.4f}%'
             .format(epoch+1, args.max_epochs, meter_loss.avg, meter_rpn_loss_cls.avg, meter_rpn_loss_bbox.avg, meter_rpn_acc.avg))
    return float("{:.4f}".format(meter_loss.avg))

def test(args, model, device, test_loader, output_file, region_output_file):
    """Test the model"""
    write_post = open_or_fd(output_file, "wb")                                  
    fid = open(region_output_file, "w")
    model.eval()
    with torch.no_grad():
        total_step = len(test_loader)
        for batch_idx, (utt_ids, act_lens, data, target) in enumerate(test_loader):
            act_lens, data, target = act_lens.to(device), data.to(device), target.to(device)
            target = target.reshape(target.size(0), 1, target.size(1)).float()
            # Forward pass
            batch_size = data.shape[0]
            max_lens = data.shape[1]
            num_anchors_per_frame = args.num_anchor
            num_classes = args.output_dim
            outputs = model(0, data, act_lens, target, 100)
            rois, rpn_cls_score, anchors_per_utt = outputs
            rpn_cls_prob = F.softmax(rpn_cls_score, dim=2)
            disable_indexes = get_out_utt_boxes(anchors_per_utt, act_lens, batch_size)
            rpn_cls_prob[disable_indexes] = 0
            rpn_cls_prob = rpn_cls_prob.view(batch_size, max_lens, num_anchors_per_frame, num_classes)
            rois = rois.view(batch_size, max_lens, num_anchors_per_frame, 2)
            anchors_per_utt = anchors_per_utt.view(max_lens, num_anchors_per_frame, 2)
            rpn_cls_prob, arg_max_anchor = torch.max(rpn_cls_prob, dim=2)
            max_score, arg_max_score = torch.max(rpn_cls_prob, dim=1) # get the index of each utterance
            data_write = rpn_cls_prob.cpu().numpy()
            for i in range (len(utt_ids)):
                utt_id = utt_ids[i]
                act_len = act_lens[i]
                write_mat(write_post, data_write[i,0:act_len,:], utt_id)
                fid.writelines(utt_id)
                label = target[i][0].cpu().numpy()
                fid.writelines(", %f %f %f"%(label[0],label[1],label[2])) 
                for j in range(num_classes-1):
                    best_score1 = max_score[i][1+j]
                    best_frame1 = arg_max_score[i][1+j]
                    best_anchor1 = arg_max_anchor[i][best_frame1][1+j]
                    roi1 = rois[i][best_frame1][best_anchor1] # anchor of keyword 1
                    anchor1 = anchors_per_utt[best_frame1][best_anchor1]
                    roi1=roi1.cpu().numpy()
                    anchor1 = anchor1.cpu().numpy()
                    fid.writelines(", %f %f %f, %f %f %f"%(best_score1, anchor1[0], anchor1[1], best_score1, roi1[0], roi1[1]))
                fid.writelines("\n")
    write_post.close()
    fid.close()
def main():
    args = get_args()

    device = torch.device('cuda' if args.use_cuda else 'cpu')
    torch.manual_seed(args.seed)
    if args.encoder=='gru':
        feature_extractor = GRU(input_size=args.input_dim, 
                output_size=args.hidden_dim, 
                hidden_size=args.hidden_dim, 
                num_layers=args.num_layers, 
                bias=True, batch_first=True, 
                dropout=args.dropout, 
                bidirectional=False, 
                output_layer=False)
    else:
        print("unsupported feature extractor: %s"%args.encoder)
        exit(1)

    rpn = RPN(128, args.num_anchor, args.output_dim)

    model = RPN_KWS(feature_extractor, rpn, args.output_dim).to(device)

    params = count_parameters(model)                                            
    print("Num parameters: %d, Num Flops: %d\n"%(params,0))
    
    if args.multi_gpu:
        model = nn.DataParallel(model)

    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=args.init_weight_decay)

    print("Global Config:\n {}".format(cfg))
    print("Training Arguments:\n {}".format(args))
    print("Training Model:\n {}".format(model))
    print("Training Optimizer:\n {}".format(optimizer))

    # Load previous trained model
    if args.load_model != '':
        print("=> Loading previous checkpoint to train: {}".format(args.load_model))
        checkpoint = torch.load(args.load_model)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        prev_val_loss = checkpoint['prev_val_loss']
    elif not args.train:
        sys.exit("Option --load-model should not be empty for testing.")
    else:
        print("=> No checkpoint found.")
        prev_val_loss = float('inf')

    # For training
    if args.train:
        if args.train_scp == '' or args.dev_scp == '':
            sys.exit("Options --train-scp and --dev-scp are required for training.")

        if args.save_dir == '':
            sys.exit("Option --save-dir is required to save model.")

        if not os.path.exists(args.save_dir):
            os.makedirs(args.save_dir)

        halving = 0
        best_model = args.load_model
        kwargs = {'num_workers': 3, 'pin_memory': True} if args.use_cuda else {}

        # Training data loader
        train_set = StreamingTorchDataset(args.train_scp, ["kaldi_reader", "raw_list_reader"], args.left_context, args.right_context)
        train_loader = torch.utils.data.DataLoader(
            dataset=train_set,
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=args.num_workers,
            collate_fn=collate_fn)

        # Dev data loader
        dev_set = StreamingTorchDataset(args.dev_scp,["kaldi_reader", "raw_list_reader"], args.left_context, args.right_context)
        dev_loader = torch.utils.data.DataLoader(
            dataset=dev_set,
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=args.num_workers,
            collate_fn=collate_fn)

        for epoch in range(args.max_epochs):
            cur_tr_loss = train(args, model, device, train_loader,optimizer, epoch)
            cur_val_loss = validate(args, model, device, dev_loader, epoch)
            rel_impr = (prev_val_loss - cur_val_loss) / prev_val_loss

            model_name = 'nnet_epoch' + str(epoch+1) + '_lr' \
                        + str(args.learning_rate) + '_tr' + str(cur_tr_loss) \
                        + '_cv' + str(cur_val_loss) + '.ckpt'
            model_path = args.save_dir + '/' + model_name

            if cur_val_loss < prev_val_loss:

                prev_val_loss = cur_val_loss
                torch.save({
                    'prev_val_loss': prev_val_loss,
                    'model': model.state_dict(),
                    'optimizer': optimizer.state_dict()
                }, model_path)
                best_model = model_path

                print("Model {} accepted. Time: {}".format(model_name,
                                                           datetime.datetime.now()))

            else:
                print ("Model {} rejected. Time: {}".format(model_name,
                                                            datetime.datetime.now()))
                if best_model != '':
                    print("=> Loading best checkpoint: {}".format(best_model))
                    checkpoint = torch.load(best_model)
                    model.load_state_dict(checkpoint['model'])
                    optimizer.load_state_dict(checkpoint['optimizer'])
                    prev_val_loss = checkpoint['prev_val_loss']
                else:
                    sys.exit("Error training neural network.")

            # Stopping training criterion
            if halving and rel_impr < args.end_halving_impr:
                if epoch < args.min_epochs:
                    print("We were supposed to finish, but we continue as min_epochs"
                          .format(args.min_epochs))
                    continue
                else:
                    print("Finished, too small relative improvement {}".format(rel_impr))
                    break

            # Start halving when improvement is low
            if rel_impr < args.start_halving_impr:
                halving = 1

            if halving:
                adjust_learning_rate(args, optimizer)
                print("Halving learning rate to {}".format(args.learning_rate))

        if best_model != args.load_model:
            final_model = args.save_dir + "/final.mdl"
            shutil.copyfile(best_model, final_model)
            print("Succeeded training the neural network: {}/final.mdl"
                  .format(args.save_dir))
        else:
            sys.exit("Error training neural network.")
    # For testing
    if args.test:
        # Test data loader
        if args.test_scp == '' or args.output_file == '':
            sys.exit("Options --test-scp and --output-file are required for testing")
        test_set = StreamingTorchDataset(args.test_scp,["kaldi_reader", "raw_list_reader"], args.left_context, args.right_context)
        test_loader = torch.utils.data.DataLoader(
            dataset=test_set,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.num_workers,
            collate_fn=collate_fn)
        test(args, model, device, test_loader, args.output_file, args.region_output_file) 


if __name__ == '__main__':
    main()