main.py

# -*- coding: utf-8 -*-
"""
Created on Wed Jul  8 05:37:22 2020

@author: Austin Hsu
"""

from src.solver import OnOffsetSolver
from src.trainer import Trainer
import argparse
import os

def train(hparams):
    
    trainer = Trainer(hparams)
    solver = OnOffsetSolver(hparams)
    
    trainer.fit(solver)
    
    trainer.test(solver)
    
    return

def test(hparams):
    
    trainer = Trainer(hparams)
    solver = OnOffsetSolver(hparams)

    solver.load_from_checkpoint(os.path.join(hparams.save_path, hparams.checkpoint_name), use_gpu=hparams.use_gpu)
    
    trainer.test(solver)
    
    return

if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    
    # --- Solver Args ---
    parser.add_argument('--model_type', help='Model used', default='PyramidNet_ShakeDrop', type=str)
    parser.add_argument('--loss_type', help='Loss used', default='VAT', type=str)
    parser.add_argument('--dataset1', help='Supervised dataset', default='TONAS', type=str)
    parser.add_argument('--dataset2', help='Semi-supervised dataset', default='Pop_Rhythm', type=str)
    parser.add_argument('--dataset3', help='Instrumental dataset', default='Pop_Rhythm_Instrumental', type=str)
    parser.add_argument('--dataset4', help='Validation dataset', default='None', type=str)
    parser.add_argument('--dataset5', help='Test dataset', default='DALI', type=str)
    parser.add_argument('--mix_ratio', help='Ratio of instrumental dataset to mix with supervised dataset', default=0.5, type=float)
    parser.add_argument('--lt_on_smooth', help='Use longtail function to smoothen onset', action='store_true')
    parser.add_argument('--lt_off_smooth', help='Use longtail function to smoothen offset', action='store_true')
    parser.add_argument('--meta_path', help='Datapath to meta files', default='./meta/', type=str)
    parser.add_argument('--data_path', help='Datapath to wav files', default='../data/', type=str)
    parser.add_argument('--lr', help='Learning rate', default=0.0001, type=float)
    parser.add_argument('--lr_warmup', help='Warmup steps for learning rate', default=40000, type=int)
    parser.add_argument('--max_steps', help='Max training steps', default=240000, type=int)
    parser.add_argument('--max_epoch', help='Max training epoch', default=20, type=int)
    parser.add_argument('--se', help='Re-train same sample for how many epochs', default=2, type=int)
    parser.add_argument('--num_feat', help='Numbers of channel dimension', default=9, type=int)
    parser.add_argument('--k', help='Window side from one side (Window size=2*k+1)', default=9, type=int)
    parser.add_argument('--batch_size', help='Batch size', default=64, type=int)
    parser.add_argument('--num_workers', help='Num workers', default=1, type=int)
    parser.add_argument('--use_ground_truth', help='Use ground truth pitch file, otherwise use patch CNN for pitch extraction', action='store_true')
    parser.add_argument('--shuffle', help='Shuffle dataset or not', action='store_true')
    parser.add_argument('--pin_memory', help='Pin memory', action='store_true')
    parser.add_argument('--use_cp', help='Use CuPy instead of Numpy for data preprocessing', action='store_true')

    # --- Trainer Args ---
    parser.add_argument('--exp_name', help='Name of current experiment', default='Train_0', type=str)
    parser.add_argument('--log_path', help='Path to save log files', default='./log/', type=str)
    parser.add_argument('--save_path', help='Path to save weight', default='./checkpoints/', type=str)
    parser.add_argument('--project', help='Wandb project name', default='note_segmentation', type=str)
    parser.add_argument('--entity', help='Wandb user name', default='austinhsu', type=str)
    parser.add_argument('--checkpoint_name', help='Checkpoint .pt file name', default='epoch=0', type=str)
    parser.add_argument('--amp_level', help='Amp level for mixed precision training', default='O0', type=str)
    parser.add_argument('--accumulate_grad_batches', help='Gradient accumulation', default=1, type=int)
    parser.add_argument('--skip_val', help='Skip validation loop, use train loss for best model', action='store_true')
    parser.add_argument('--use_gpu', help='Use gpu for training', action='store_true')
    parser.add_argument('--test_no_offset', help='Test without offset_ratio, recommended for offset-unstable datasets (like DALI)', action='store_true')
    parser.add_argument('--train', help='Run in train mode', action='store_true')
    parser.add_argument('--test', help='Run in test mode', action='store_true')

    #parser.set_defaults(
    #    # --- Solver Args ---
    #    use_ground_truth=False,
    # 	 shuffle=True,
    # 	 pin_memory=True,
    # 	 use_cp=True,

    #    # --- Trainer Args ---
    #    skip_val=False,
    #    use_gpu=True,
    #    train=True,
    #    test=False,
    #)

    hyperparams = parser.parse_args()

    if hyperparams.train:
        train(hyperparams)

    if hyperparams.test:
        test(hyperparams)