def train(args): from train_utils import ModelTrainer device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Prepare Data train_dataset = Chime_Dataset('tr', args) train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, pin_memory=True, collate_fn=lambda x: Chime_Collate(x), num_workers=args.num_workers) # Prepare model if args.model_type == 'BLSTM': model = BLSTMMaskEstimator() model_save_dir = os.path.join(args.data_dir, 'BLSTM_model') mkdir_p(model_save_dir) elif args.model_type == 'FW': model = SimpleFWMaskEstimator() model_save_dir = os.path.join(args.data_dir, 'FW_model') mkdir_p(model_save_dir) else: raise ValueError('Unknown model type. Possible are "BLSTM" and "FW"') criterion = torch.nn.BCELoss() optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=1e-4) trainer = ModelTrainer(model, train_loader, criterion, optimizer, args, device) trainer.train(args.num_epochs)
def test(args): from test_utils import ModelTest device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Prepare Data test_dataset = Chime_Dataset('dt', args) test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=True, pin_memory=True, collate_fn=lambda x: Chime_Collate(x), num_workers=args.num_workers) # Prepare model if args.model_type == 'BLSTM': model = BLSTMMaskEstimator() model_save_dir = os.path.join(args.data_dir, 'BLSTM_model') mkdir_p(model_save_dir) elif args.model_type == 'FW': model = SimpleFWMaskEstimator() model_save_dir = os.path.join(args.data_dir, 'FW_model') mkdir_p(model_save_dir) else: raise ValueError('Unknown model type. Possible are "BLSTM" and "FW"') criterion = torch.nn.BCELoss() tester = ModelTest(model, test_loader, criterion, args, device) tester.test()
parser = argparse.ArgumentParser(description='NN GEV beamforming') parser.add_argument('model', help='Trained model file') parser.add_argument('model_type', help='Type of model (BLSTM or FW)') parser.add_argument('--gpu', '-g', default=-1, type=int, help='GPU ID (negative value indicates CPU)') parser.add_argument('data_directory', help='data experiment directory') parser.add_argument('exNum', help='Experiment order') args = parser.parse_args() # Prepare model if args.model_type == 'BLSTM': model = BLSTMMaskEstimator() elif args.model_type == 'FW': model = SimpleFWMaskEstimator() else: raise ValueError('Unknown model type. Possible are "BLSTM" and "FW"') serializers.load_hdf5(args.model, model) print("data type of 'model'", type(model)) if args.gpu >= 0: cuda.get_device(args.gpu).use() model.to_gpu() xp = np if args.gpu < 0 else cuda.cupy # def single_noise(): # audio_data = get_audio_nochime('new_dataset/2m/2m_pub_new', ch_range=range(1, 9), fs=49000) # # audio_data = get_audio_nochime('new_dataset/new_audio/AUDIO_RECORDING', ch_range=range(1, 9), fs=49000)
help='The directory where the enhanced wav files will ' 'be stored.') parser.add_argument('model', help='Trained model file') parser.add_argument('model_type', help='Type of model (BLSTM or FW)') parser.add_argument('beamformer', help='Type of beamformer (GEV, MVDR of ALL)') parser.add_argument('--gpu', '-g', default=-1, type=int, help='GPU ID (negative value indicates CPU)') args = parser.parse_args() # Prepare model if args.model_type == 'BLSTM': model = BLSTMMaskEstimator() elif args.model_type == 'FW': model = SimpleFWMaskEstimator() else: raise ValueError('Unknown model type. Possible are "BLSTM" and "FW"') if args.beamformer == 'GEV': beamformers=['gev'] elif args.beamformer == 'MVDR': beamformers=['mvdr'] elif args.beamformer == 'ALL': beamformers=['gev','mvdr'] print("Selected beamformers:",' '.join(beamformers).upper()) serializers.load_hdf5(args.model, model)
if args.chime_dir != '': log.info('Preparing training data and storing it in {}'.format( args.data_dir)) prepare_training_data(args.chime_dir, args.data_dir) flists = dict() for stage in ['tr', 'dt']: with open(os.path.join(args.data_dir, 'flist_{}.json'.format(stage))) as fid: flists[stage] = json.load(fid) log.debug('Loaded file lists') # Prepare model if args.model_type == 'BLSTM': model = BLSTMMaskEstimator() model_save_dir = os.path.join(args.data_dir, 'BLSTM_model') mkdir_p(model_save_dir) elif args.model_type == 'FW': model = SimpleFWMaskEstimator() model_save_dir = os.path.join(args.data_dir, 'FW_model') mkdir_p(model_save_dir) else: raise ValueError('Unknown model type. Possible are "BLSTM" and "FW"') if args.gpu >= 0: cuda.get_device(args.gpu).use() model.to_gpu() xp = np if args.gpu < 0 else cuda.cupy log.debug('Prepared model')
from nn_models import BLSTMMaskEstimator, SimpleFWMaskEstimator parser = argparse.ArgumentParser(description='NN GEV beamforming') parser.add_argument('model', help='Trained model file') parser.add_argument('model_type', help='Type of model (BLSTM or FW)') parser.add_argument('--gpu', '-g', default=-1, type=int, help='GPU ID (negative value indicates CPU)') parser.add_argument('experiments', help='Number of current experiment') args = parser.parse_args() # Prepare model if args.model_type == 'BLSTM': model = BLSTMMaskEstimator() elif args.model_type == 'FW': model = SimpleFWMaskEstimator() else: raise ValueError('Unknown model type. Possible are "BLSTM" and "FW"') """ PS: feedback loop beamforming experiment, using beamforming result + specified channel. the beamforming as the channel """ # load nnet model serializers.load_hdf5(args.model, model) t_io = 0 t_net = 0 t_beamform = 0
parser.add_argument('chime_dir', help='Base directory of the CHiME challenge.') parser.add_argument('output_dir', help='The directory where the enhanced wav files will ' 'be stored.') parser.add_argument('model', help='Trained model file') parser.add_argument('model_type', help='Type of model (BLSTM or FW)') parser.add_argument('--gpu', '-g', default=-1, type=int, help='GPU ID (negative value indicates CPU)') args = parser.parse_args() # Prepare model if args.model_type == 'BLSTM': model = BLSTMMaskEstimator() elif args.model_type == 'FW': model = SimpleFWMaskEstimator() else: raise ValueError('Unknown model type. Possible are "BLSTM" and "FW"') serializers.load_hdf5(args.model, model) if args.gpu >= 0: cuda.get_device(args.gpu).use() model.to_gpu() xp = np if args.gpu < 0 else cuda.cupy stage = args.flist[:2] scenario = args.flist.split('_')[-1] # CHiME data handling
if args.chime_dir != '': log.info( 'Preparing training data and storing it in {}'.format( args.data_dir)) prepare_training_data(args.chime_dir, args.data_dir) flists = dict() for stage in ['tr', 'dt']: with open( os.path.join(args.data_dir, 'flist_{}.json'.format(stage))) as fid: flists[stage] = json.load(fid) log.debug('Loaded file lists') # Prepare model if args.model_type == 'BLSTM': model = BLSTMMaskEstimator() model_save_dir = os.path.join(args.data_dir, 'BLSTM_model') mkdir_p(model_save_dir) elif args.model_type == 'FW': model = SimpleFWMaskEstimator() model_save_dir = os.path.join(args.data_dir, 'FW_model') mkdir_p(model_save_dir) else: raise ValueError('Unknown model type. Possible are "BLSTM" and "FW"') if args.gpu >= 0: cuda.get_device(args.gpu).use() model.to_gpu() xp = np if args.gpu < 0 else cuda.cupy log.debug('Prepared model')