def main(args): # Construct Solver # data tr_dataset = AudioDataset(args.train_json, sample_rate=args.sample_rate, segment_length=args.segment_length) cv_dataset = AudioDataset( args.valid_json, sample_rate=args.sample_rate, segment_length=args.segment_length, ) tr_loader = AudioDataLoader(tr_dataset, batch_size=args.batch_size, shuffle=args.shuffle, num_workers=args.num_workers) cv_loader = AudioDataLoader(cv_dataset, batch_size=args.batch_size, num_workers=0) data = {'tr_loader': tr_loader, 'cv_loader': cv_loader} # model model = ConvTasNet(args.N, args.L, args.B, args.H, args.P, args.X, args.R, args.C, norm_type=args.norm_type, causal=args.causal, mask_nonlinear=args.mask_nonlinear) print(model) if args.use_cuda: model = torch.nn.DataParallel(model) model.cuda() # optimizer lr = args.lr / args.batch_per_step if args.optimizer == 'sgd': optimizier = torch.optim.SGD(model.parameters(), lr=lr, momentum=args.momentum, weight_decay=args.l2) elif args.optimizer == 'adam': optimizier = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=args.l2) else: print("Not support optimizer") return # solver solver = Solver(data, model, optimizier, args) solver.train()
def main(args): # Construct Solver # data tr_dataset = AudioDataset(args.train_dir, args.batch_size, sample_rate=args.sample_rate, segment=args.segment) cv_dataset = AudioDataset( args.valid_dir, batch_size=1, # 1 -> use less GPU memory to do cv sample_rate=args.sample_rate, segment=-1, cv_maxlen=args.cv_maxlen) # -1 -> use full audio tr_loader = AudioDataLoader(tr_dataset, batch_size=1, shuffle=args.shuffle, num_workers=4) cv_loader = AudioDataLoader(cv_dataset, batch_size=1, num_workers=4, pin_memory=True) data = {'tr_loader': tr_loader, 'cv_loader': cv_loader} # model model = ConvTasNet(args.N, args.L, args.B, args.H, args.P, args.X, args.R, args.C, norm_type=args.norm_type, causal=args.causal, mask_nonlinear=args.mask_nonlinear) if args.use_cuda: model = torch.nn.DataParallel(model) model.cuda() # optimizer if args.optimizer == 'sgd': optimizier = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.l2) elif args.optimizer == 'adam': optimizier = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2) else: print("Not support optimizer") return # solver solver = Solver(data, model, optimizier, args) solver.train()
def main(train_dir,batch_size,sample_rate, segment,valid_dir,cv_maxlen,shuffle,num_workers,N, L, B, H, P, X, R, C,norm_type, causal, mask_nonlinear,use_cuda,optimizer,lr,momentum,l2): # Construct Solver # data tr_dataset = AudioDataset(train_dir, batch_size, sample_rate=sample_rate, segment=segment) cv_dataset = AudioDataset(valid_dir, batch_size=1, # 1 -> use less GPU memory to do cv sample_rate=sample_rate, segment=-1, cv_maxlen=cv_maxlen) # -1 -> use full audio tr_loader = AudioDataLoader(tr_dataset, batch_size=1, shuffle=shuffle, num_workers=num_workers) cv_loader = AudioDataLoader(cv_dataset, batch_size=1, num_workers=0) data = {'tr_loader': tr_loader, 'cv_loader': cv_loader} # model model = ConvTasNet(N, L, B, H, P, X, R, C, norm_type=norm_type, causal=causal, mask_nonlinear=mask_nonlinear) print(model) if use_cuda: model = torch.nn.DataParallel(model) model.cuda() # optimizer if optimizer == 'sgd': optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum, weight_decay=l2) elif optimizer == 'adam': #fatemeh: change optimizier to optimizer optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=l2) else: print("Not support optimizer") return # solver solver = Solver(data, model, optimizer, use_cuda,epochs,half_lr,early_stop,max_norm,save_folder,checkpoint,continue_from,model_path,print_freq,visdom,visdom_epoch,visdom_id) solver.train()
def main(args): # Construct Solver # data if args.continue_from == '': return ev_dataset = EvalAllDataset(args.train_dir, args.mix_json, args.batch_size, sample_rate=args.sample_rate) ev_loader = EvalAllDataLoader(ev_dataset, batch_size=1, num_workers=args.num_workers) data = {'tr_loader': None, 'ev_loader': ev_loader} # SEP model sep_model = ConvTasNet(args.N, args.L, args.B, args.H, args.P, args.X, args.R, args.C, norm_type=args.norm_type, causal=args.causal, mask_nonlinear=args.mask_nonlinear) # ASR model asr_model = AttentionModel(args.NUM_HIDDEN_NODES, args.NUM_ENC_LAYERS, args.NUM_CLASSES) #print(model) if args.use_cuda: sep_model = torch.nn.DataParallel(sep_model) asr_model = torch.nn.DataParallel(asr_model) sep_model.cuda() asr_model.cuda() # optimizer if args.optimizer == 'sgd': sep_optimizier = torch.optim.SGD(sep_model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.l2) asr_optimizier = torch.optim.SGD(asr_model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.l2) elif args.optimizer == 'adam': sep_optimizier = torch.optim.Adam(sep_model.parameters(), lr=args.lr, weight_decay=args.l2) asr_optimizier = torch.optim.Adam(asr_model.parameters(), lr=args.lr, weight_decay=args.l2) else: print("Not support optimizer") return # solver solver = Solver(data, sep_model, asr_model, sep_optimizier, asr_optimizier, args, DEVICE, ev=True) solver.eval(args.EOS_ID)