test_run_model.py

import numpy as np
import argparse
import logging
import os
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from model.utils import move_data_to_device

from params import train_config
from data.utils import create_folder, create_logging, get_filename
from data.data_loader import AudioDataset, TrainSampler, EvaluateSampler, collate_fn
#from model.models import Transfer_Cnn14
#from model.losses import get_loss_func
#from model.evaluate import Eva

def train(args):

    # Parameters
    dataset_dir = args.dataset_dir
    workspace = args.workspace
    holdout_fold = args.holdout_fold
    model_type = args.model_type
    pretrained_checkpoint_path = args.pretrained_checkpoint_path
    freeze_base = args.freeze_base
    loss_type = args.loss_type
    augmentation = args.augmentation
    learning_rate = args.learning_rate
    batch_size = args.batch_size
    resume_iteration = args.resume_iteration
    stop_iteration = args.stop_iteration
    device = 'cuda' if (args.cuda and torch.cuda.is_available()) else 'cpu'
    filename = args.filename
    num_workers = 1
    #loss_func = get_loss_func(loss_type)
    pretrain = True if pretrained_checkpoint_path else False

    hdf5_path = os.path.join(workspace, 'features', 'waveform.h5')

    checkpoints_dir = os.path.join(workspace, 'checkpoints', filename,
        'holdout_fold={}'.format(holdout_fold), model_type, 'pretrain={}'.format(pretrain),
        'loss_type={}'.format(loss_type), 'augmentation={}'.format(augmentation),
        'batch_size={}'.format(batch_size), 'freeze_base={}'.format(freeze_base)
    )
    create_folder(checkpoints_dir)

    # Model = eval(model_type) # This could be Model = Transfer_Cnn14() in our case, however, here for easy implementation, we will still use this.
    # model = Model(train_config.sample_rate, train_config.window_size, train_config.hop_size, train_config.mel_bins,
    # train_config.fmin, train_config.fmax, train_config.classes_num, train_config.freeze_base)

    # if pretrain:
    #     print("Load pretrained model from {}".format(pretrained_checkpoint_path))
    #     #model.load_from_pretrain(pretrained_checkpoint_path)

    # if resume_iteration:
    #     resume_checkpoint_path = os.path.join(checkpoints_dir, '{}_iterations.pth'.format(resume_iteration))
    #     print("Load resume model from {}".format(resume_checkpoint_path))
    #     resume_checkpoint = torch.load(resume_checkpoint_path)
    #     model.load_state_dict(resume_checkpoint['model'])
    #     iteration = resume_checkpoint['iteration']
    # else:
    #     iteration = 0

    # Data
    dataset = AudioDataset()

    # Generator
    train_sampler = TrainSampler(
        hdf5_path=hdf5_path,
        holdout_fold=holdout_fold,
        batch_size=batch_size
    )



    validate_sampler = EvaluateSampler(
        hdf5_path=hdf5_path,
        holdout_fold=holdout_fold,
        batch_size=batch_size
    )

    print('line94')
    # Data Loader
    train_loader = torch.utils.data.DataLoader(dataset=dataset,
        batch_sampler=train_sampler, collate_fn=collate_fn,
        num_workers=num_workers, pin_memory=True
    )


    print('line100')

    validate_loader = torch.utils.data.DataLoader(dataset=dataset,
        batch_sampler=validate_sampler, collate_fn=collate_fn,
        num_workers=num_workers, pin_memory=True
    )

    # if 'cuda' in device:
    #     model.to(device)

    # # Optimizer
    # optimizer = optim.Adam(model.parameters(), lr=learning_rate, betas=(0.9, 0.999),
    #     eps=1e-08, weight_decay=0., amsgrad=True
    # )

    # # Evaluator
    # evaluator = Eva(model=model)

    # train_begin_time = time.time()
    # print(train_begin_time)

    # Train
    print('Start Test')
    for batch_data_dict in train_loader:
        
        for key in batch_data_dict.keys():
            batch_data_dict[key] = move_data_to_device(batch_data_dict[key], device)

        #print(batch_data_dict['waveform'])
        
        print(np.array([element.decode("utf-8") for element in batch_data_dict['caption']]))
        

        #print(batch_data_dict['fold_num'])

    # for batch_data_dict in train_loader:
        
    #     # Evaluate
    #     if iteration % 100 == 0 and iteration > 0:
    #         if resume_iteration > 0 and iteration == resume_iteration:
    #             pass
    #         else:
    #             print("-----------------------------------------------")
    #             print("Iteration: {}".format(iteration))

    #             train_fin_time = time.time()
    #             statistics =  evaluator.evaluate(validate_loader)
    #             print("Validate accuracy: {:.3f}".format(statistics['accuracy']))

    #             train_time = train_fin_time - train_begin_time
    #             validate_time = time.time() - train_fin_time

    #             '''
    #             logging.info(
    #                 "Train time: {:.3f} s, validate time: {:.3f} s".format(train_time, validate_time)
    #             )'''

    #             train_begin_time = time.time()
    #     # Save
    #     '''
    #     if iteration % 2000 == 0 or iteration > 0:
    #         checkpoint = {
    #             'iteration': iteration,
    #             'model': model.module.state_dict()
    #         }

    #         checkpoint_path = os.path.join(checkpoints_dir, '{}_iterations.pth'.format(iteration))
    #         torch.save(checkpoint, checkpoint_path)
    #         print('Model saved to {}'.format(checkpoint_path))'''
            
    #     # Move data to GPU
    #     for key in batch_data_dict.keys():
    #         batch_data_dict[key] = move_data_to_device(batch_data_dict[key], device)

    #     # Train
    #     model.train()

    #     batch_output_dict = model(batch_data_dict['waveform'], None)
        
    #     batch_targets_dict = {'target': batch_data_dict['target']}

    #     # loss
    #     loss = loss_func(batch_output_dict, batch_targets_dict)
    #     if iteration % 100 == 0 and iteration > 0:
    #         print(iteration, loss)

    #     # Backward
    #     optimizer.zero_grad()
    #     loss.backward()
    #     optimizer.step()

    #     # Stop
    #     if iteration == stop_iteration:
    #         break

    #     iteration += 1


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="")
    subparsers = parser.add_subparsers(dest='mode')

    # Train
    parser_train = subparsers.add_parser("train")
    parser_train.add_argument("--dataset_dir", type=str, required=True, help='Directory of dataset. ')
    parser_train.add_argument("--workspace", type=str, required=True, help='Directory of your workspace. ')
    parser_train.add_argument("--holdout_fold", type=str, choices=['1', '2', '3', '4', '5'], required=True)
    parser_train.add_argument("--model_type", type=str, required=True)
    parser_train.add_argument("--pretrained_checkpoint_path", type=str)
    parser_train.add_argument("--freeze_base", action='store_true', default=False)
    parser_train.add_argument("--loss_type", type=str, required=True)
    parser_train.add_argument("--augmentation", type=str, choices=['none', 'mixup'], required=True) # for easy implementation, I set it to False
    parser_train.add_argument("--learning_rate", type=float, required=True)
    parser_train.add_argument("--batch_size", type=int, required=True)
    parser_train.add_argument("--resume_iteration", type=int)
    parser_train.add_argument("--stop_iteration", type=int, required=True)
    parser_train.add_argument("--cuda", action='store_true', default=False)

    # Parse arguments
    args = parser.parse_args()
    args.filename = get_filename(__file__)

    if args.mode == 'train':
        train(args)
    else:
        raise Exception("Args mode should be train. ")