def train(config: DictConfig) -> nn.DataParallel: random.seed(config.train.seed) torch.manual_seed(config.train.seed) torch.cuda.manual_seed_all(config.train.seed) device = check_envirionment(config.train.use_cuda) if hasattr(config.train, "num_threads") and int(config.train.num_threads) > 0: torch.set_num_threads(config.train.num_threads) vocab = KsponSpeechVocabulary( f'/home/seungmin/dmount/KoSpeech/data/vocab/aihub_{config.train.output_unit}_vocabs.csv', output_unit=config.train.output_unit, ) if not config.train.resume: epoch_time_step, trainset_list, validset = split_dataset( config, config.train.transcripts_path, vocab) model = build_model(config, vocab, device) optimizer = get_optimizer(model, config) lr_scheduler = get_lr_scheduler(config, optimizer, epoch_time_step) optimizer = Optimizer(optimizer, lr_scheduler, config.train.total_steps, config.train.max_grad_norm) criterion = get_criterion(config, vocab) else: trainset_list = None validset = None model = None optimizer = None epoch_time_step = None criterion = get_criterion(config, vocab) trainer = SupervisedTrainer( optimizer=optimizer, criterion=criterion, trainset_list=trainset_list, validset=validset, num_workers=config.train.num_workers, device=device, teacher_forcing_step=config.model.teacher_forcing_step, min_teacher_forcing_ratio=config.model.min_teacher_forcing_ratio, print_every=config.train.print_every, save_result_every=config.train.save_result_every, checkpoint_every=config.train.checkpoint_every, architecture=config.model.architecture, vocab=vocab, joint_ctc_attention=config.model.joint_ctc_attention, ) model = trainer.train( model=model, batch_size=config.train.batch_size, epoch_time_step=epoch_time_step, num_epochs=config.train.num_epochs, teacher_forcing_ratio=config.model.teacher_forcing_ratio, resume=config.train.resume, ) return model
def train(config: DictConfig) -> nn.DataParallel: random.seed(config.train.seed) torch.manual_seed(config.train.seed) torch.cuda.manual_seed_all(config.train.seed) device = check_envirionment(config.train.use_cuda) if config.train.dataset == 'kspon': if config.train.output_unit == 'subword': vocab = KsponSpeechVocabulary( vocab_path=KSPONSPEECH_VOCAB_PATH, output_unit=config.train.output_unit, sp_model_path=KSPONSPEECH_SP_MODEL_PATH, ) else: vocab = KsponSpeechVocabulary( f'../../../data/vocab/aihub_{config.train.output_unit}_vocabs.csv', output_unit=config.train.output_unit, ) elif config.train.dataset == 'libri': vocab = LibriSpeechVocabulary(LIBRISPEECH_VOCAB_PATH, LIBRISPEECH_TOKENIZER_PATH) else: raise ValueError("Unsupported Dataset : {0}".format(config.train.dataset)) if not config.train.resume: epoch_time_step, trainset_list, validset = split_dataset(config, config.train.transcripts_path, vocab) model = build_model(config, vocab, device) optimizer = get_optimizer(model, config) lr_scheduler = get_lr_scheduler(config, optimizer, epoch_time_step) optimizer = Optimizer(optimizer, lr_scheduler, config.train.warmup_steps, config.train.max_grad_norm) criterion = get_criterion(config, vocab) else: trainset_list = None validset = None model = None optimizer = None epoch_time_step = None criterion = get_criterion(config, vocab) trainer = SupervisedTrainer( optimizer=optimizer, criterion=criterion, trainset_list=trainset_list, validset=validset, num_workers=config.train.num_workers, device=device, teacher_forcing_step=config.model.teacher_forcing_step, min_teacher_forcing_ratio=config.model.min_teacher_forcing_ratio, print_every=config.train.print_every, save_result_every=config.train.save_result_every, checkpoint_every=config.train.checkpoint_every, architecture=config.model.architecture, vocab=vocab, joint_ctc_attention=config.model.joint_ctc_attention, ) model = trainer.train( model=model, batch_size=config.train.batch_size, epoch_time_step=epoch_time_step, num_epochs=config.train.num_epochs, teacher_forcing_ratio=config.model.teacher_forcing_ratio, resume=config.train.resume, ) return model
def train(config: DictConfig): random.seed(config.train.seed) torch.manual_seed(config.train.seed) torch.cuda.manual_seed_all(config.train.seed) device = check_envirionment(config.train.use_cuda) if config.train.dataset == 'kspon': if config.train.output_unit == 'subword': vocab = KsponSpeechVocabulary( vocab_path='../../../data/vocab/kspon_sentencepiece.vocab', output_unit=config.train.output_unit, sp_model_path='../../../data/vocab/kspon_sentencepiece.model', ) else: vocab = KsponSpeechVocabulary( f'../../../data/vocab/aihub_{config.train.output_unit}_vocabs.csv', output_unit=config.train.output_unit) elif config.train.dataset == 'libri': vocab = LibriSpeechVocabulary('../../../data/vocab/tokenizer.vocab', '../../../data/vocab/tokenizer.model') else: raise ValueError("Unsupported Dataset : {0}".format( config.train.dataset)) if not config.train.resume: epoch_time_step, trainset_list, validset = split_dataset( config, config.train.transcripts_path, vocab) model = build_model(config, vocab, device) optimizer = get_optimizer(model, config) lr_scheduler = TriStageLRScheduler( optimizer=optimizer, init_lr=config.train.init_lr, peak_lr=config.train.peak_lr, final_lr=config.train.final_lr, init_lr_scale=config.train.init_lr_scale, final_lr_scale=config.train.final_lr_scale, warmup_steps=config.train.warmup_steps, total_steps=int(config.train.num_epochs * epoch_time_step)) optimizer = Optimizer(optimizer, lr_scheduler, config.train.warmup_steps, config.train.max_grad_norm) criterion = get_criterion(config, vocab) else: trainset_list = None validset = None model = None optimizer = None epoch_time_step = None criterion = get_criterion(config, vocab) trainer = SupervisedTrainer( optimizer=optimizer, criterion=criterion, trainset_list=trainset_list, validset=validset, num_workers=config.train.num_workers, device=device, teacher_forcing_step=config.model.teacher_forcing_step, min_teacher_forcing_ratio=config.model.min_teacher_forcing_ratio, print_every=config.train.print_every, save_result_every=config.train.save_result_every, checkpoint_every=config.train.checkpoint_every, architecture=config.model.architecture, vocab=vocab, joint_ctc_attention=config.model.joint_ctc_attention, ) model = trainer.train( model=model, batch_size=config.train.batch_size, epoch_time_step=epoch_time_step, num_epochs=config.train.num_epochs, teacher_forcing_ratio=config.model.teacher_forcing_ratio, resume=config.train.resume, ) return model