def from_pretrained(model_name_or_path, checkpoint_file='model.pt', data_name_or_path='.', archive_map=None, **kwargs): from fairseq import checkpoint_utils, file_utils if archive_map is not None: if model_name_or_path in archive_map: model_name_or_path = archive_map[model_name_or_path] if data_name_or_path is not None and data_name_or_path in archive_map: data_name_or_path = archive_map[data_name_or_path] # allow archive_map to set default arg_overrides (e.g., tokenizer, bpe) # for each model if isinstance(model_name_or_path, dict): for k, v in model_name_or_path.items(): if k == 'checkpoint_file': checkpoint_file = v elif (k != 'path' # only set kwargs that don't already have overrides and k not in kwargs): kwargs[k] = v model_name_or_path = model_name_or_path['path'] model_path = file_utils.load_archive_file(model_name_or_path) # convenience hack for loading data and BPE codes from model archive if data_name_or_path.startswith('.'): kwargs['data'] = os.path.abspath( os.path.join(model_path, data_name_or_path)) else: kwargs['data'] = file_utils.load_archive_file(data_name_or_path) for file, arg in { 'code': 'bpe_codes', 'bpecodes': 'bpe_codes', 'sentencepiece.bpe.model': 'sentencepiece_vocab', }.items(): print(model_path, file) path = os.path.join(model_path, file) if os.path.exists(path): kwargs[arg] = path if 'user_dir' in kwargs: utils.import_user_module( argparse.Namespace(user_dir=kwargs['user_dir'])) models, args, task = checkpoint_utils.load_model_ensemble_and_task( [ os.path.join(model_path, cpt) for cpt in checkpoint_file.split(os.pathsep) ], arg_overrides=kwargs, ) return { 'args': args, 'task': task, 'models': models, }
def setUp(self, args): if args.task == "speech_to_speech": args.normalize_waveform = False self.task = tasks.setup_task(args) self.pre_tokenizer = self.task.build_tokenizer(args) self.bpe_tokenizer = self.task.build_bpe(args) try: self.src_dict = self.task.src_dict except Exception: self.src_dict = None ensemble, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task( [args.path], arg_overrides=ast.literal_eval(args.model_overrides), task=self.task, strict=False, ) self.model = ensemble[0] if self.use_cuda: self.model.cuda() # criterion.cuda() self.model.eval() self.generator = self.task.build_generator( [self.model], args, )
def from_pretrained( model_name_or_path, checkpoint_file="model.pt", data_name_or_path=".", archive_map=None, **kwargs ): from fairseq import checkpoint_utils, file_utils if archive_map is not None: if model_name_or_path in archive_map: model_name_or_path = archive_map[model_name_or_path] if data_name_or_path is not None and data_name_or_path in archive_map: data_name_or_path = archive_map[data_name_or_path] # allow archive_map to set default arg_overrides (e.g., tokenizer, bpe) # for each model if isinstance(model_name_or_path, dict): for k, v in model_name_or_path.items(): if k == "checkpoint_file": checkpoint_file = v elif ( k != "path" # only set kwargs that don't already have overrides and k not in kwargs ): kwargs[k] = v model_name_or_path = model_name_or_path["path"] model_path = file_utils.load_archive_file(model_name_or_path) # convenience hack for loading data and BPE codes from model archive if data_name_or_path.startswith("."): kwargs["data"] = os.path.abspath(os.path.join(model_path, data_name_or_path)) else: kwargs["data"] = file_utils.load_archive_file(data_name_or_path) for file, arg in { "code": "bpe_codes", "bpecodes": "bpe_codes", "sentencepiece.bpe.model": "sentencepiece_model", "merges.txt": "bpe_merges", "vocab.json": "bpe_vocab", }.items(): path = os.path.join(model_path, file) if os.path.exists(path): kwargs[arg] = path if "user_dir" in kwargs: utils.import_user_module(argparse.Namespace(user_dir=kwargs["user_dir"])) models, args, task = checkpoint_utils.load_model_ensemble_and_task( [os.path.join(model_path, cpt) for cpt in checkpoint_file.split(os.pathsep)], arg_overrides=kwargs, ) return { "args": args, "task": task, "models": models, }
def main(args): logging.basicConfig( format='%(asctime)s | %(levelname)s | %(name)s | %(message)s', datefmt='%Y-%m-%d %H:%M:%S', level=logging.INFO, stream=sys.stdout, ) logger = logging.getLogger('fairnr_cli.extract') logger.info(args) use_cuda = torch.cuda.is_available() and not args.cpu models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [args.path], suffix=getattr(args, "checkpoint_suffix", "")) model = models[0] if use_cuda: model.cuda() if args.format == 'mc_mesh': plydata = model.encoder.export_surfaces( model.field, th=args.mc_threshold, bits=2 * args.mc_num_samples_per_halfvoxel) elif args.format == 'voxel_center': plydata = model.encoder.export_voxels(False) elif args.format == 'voxel_mesh': plydata = model.encoder.export_voxels(True) else: raise NotImplementedError # write to ply file. if not os.path.exists(args.output): os.makedirs(args.output) plydata.text = args.savetext plydata.write(open(os.path.join(args.output, args.name + '.ply'), 'wb'))
def cli_main(): parser = argparse.ArgumentParser( description= 'Extract geometry from a trained model (only for learnable embeddings).' ) parser.add_argument('--path', type=str, required=True) parser.add_argument('--output', type=str, required=True) parser.add_argument('--name', type=str, default='sparsevoxel') parser.add_argument('--user-dir', default='fairnr') args = options.parse_args_and_arch(parser) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [args.path], suffix=getattr(args, "checkpoint_suffix", ""), ) model = models[0] voxel_idx, voxel_pts = model.encoder.extract_voxels() # write to ply file. points = [(voxel_pts[k, 0], voxel_pts[k, 1], voxel_pts[k, 2], voxel_idx[k]) for k in range(voxel_idx.size(0))] vertex = np.array(points, dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('quality', 'f4')]) PlyData([PlyElement.describe(vertex, 'vertex')], text=True).write(os.path.join(args.output, args.name + '.ply'))
def __init__( self, language, temperature=1.0, model_dir=telescope_cache_folder() + "m39v1", **kwargs, ): """ model_dir should contain: 1) checkpoint.pt: the fairseq model 2) spm.model: the sentencepiece model 3) dict.src.txt: the fairseq source dictionary 4) dict.tgt.txt: the fairseq target dictionary (likely a copy of the source) lang: ISO 639-1 Code (e.g. "en"). Must be a language compatable with the model. """ if not os.path.isdir(model_dir): download_file_maybe_extract( url="http://data.statmt.org/prism/m39v1.tar", directory=telescope_cache_folder(), ) self.sp = spm.SentencePieceProcessor() self.sp.Load(model_dir + "/spm.model") self.lang = language self.temperature = temperature # this prints things and I can't figure out how to disable it with open(os.devnull, "w") as sys.stdout: ( self.models, self.args, self.task, ) = checkpoint_utils.load_model_ensemble_and_task( [ model_dir + "/checkpoint.pt", ], arg_overrides=dict(data=model_dir + "/"), ) sys.stdout = sys.__stdout__ self.use_cuda = torch.cuda.is_available() self.generator = SequenceScorer(self.task.target_dictionary, temperature=temperature) for model in self.models: if self.use_cuda: model.cuda() model.make_generation_fast_( beamable_mm_beam_size=None, need_attn=False, ) # hash model self.model_hash = hash_model(model_dir) if not self.language_support(language): raise Exception(f"{language} is not supported by {self.name}.")
def main(args): assert (args.dump_features or args.dump_waveforms or args.dump_attentions or args.dump_eos_probs or args.dump_plots) if args.max_tokens is None and args.batch_size is None: args.max_tokens = 8000 logger.info(args) use_cuda = torch.cuda.is_available() and not args.cpu task = tasks.setup_task(args) models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task( [args.path], task=task, arg_overrides=ast.literal_eval(args.model_overrides), ) model = models[0].cuda() if use_cuda else models[0] # use the original n_frames_per_step task.args.n_frames_per_step = saved_cfg.task.n_frames_per_step task.load_dataset(args.gen_subset, task_cfg=saved_cfg.task) data_cfg = task.data_cfg sample_rate = data_cfg.config.get("features", {}).get("sample_rate", 22050) resample_fn = { False: lambda x: x, True: lambda x: torchaudio.sox_effects.apply_effects_tensor( x.detach().cpu().unsqueeze(0), sample_rate, [['rate', str(args.output_sample_rate)]])[0].squeeze(0) }.get(args.output_sample_rate != sample_rate) if args.output_sample_rate != sample_rate: logger.info(f"resampling to {args.output_sample_rate}Hz") generator = task.build_generator([model], args) itr = task.get_batch_iterator( dataset=task.dataset(args.gen_subset), max_tokens=args.max_tokens, max_sentences=args.batch_size, max_positions=(sys.maxsize, sys.maxsize), ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test, required_batch_size_multiple=args.required_batch_size_multiple, num_shards=args.num_shards, shard_id=args.shard_id, num_workers=args.num_workers, data_buffer_size=args.data_buffer_size, ).next_epoch_itr(shuffle=False) Path(args.results_path).mkdir(exist_ok=True, parents=True) is_na_model = getattr(model, "NON_AUTOREGRESSIVE", False) dataset = task.dataset(args.gen_subset) vocoder = task.args.vocoder with progress_bar.build_progress_bar(args, itr) as t: for sample in t: sample = utils.move_to_cuda(sample) if use_cuda else sample hypos = generator.generate(model, sample, has_targ=args.dump_target) for result in postprocess_results(dataset, sample, hypos, resample_fn, args.dump_target): dump_result(is_na_model, args, vocoder, *result)
def __init__(self, model_dir, lang, temperature): ''' model_dir should contain: 1) checkpoint.pt: the fairseq model 2) spm.model: the sentencepiece model 3) dict.src.txt: the fairseq source dictionary 4) dict.tgt.txt: the fairseq target dictionary (likely a copy of the source) lang: ISO 639-1 Code (e.g. "en"). Must be a language compatable with the model. ''' self.sp = spm.SentencePieceProcessor() self.sp.Load(model_dir + '/spm.model') self.lang = lang self.temperature = temperature # this prints things and I can't figure out how to disable it sys.stdout = open(os.devnull, 'w') self.models, self.args, self.task = checkpoint_utils.load_model_ensemble_and_task( [ model_dir + '/checkpoint.pt', ], arg_overrides=dict(data=model_dir + '/'), ) sys.stdout = sys.__stdout__ self.use_cuda = torch.cuda.is_available() self.generator = SequenceScorer(self.task.target_dictionary, temperature=temperature) for model in self.models: if self.use_cuda: model.cuda() model.make_generation_fast_( beamable_mm_beam_size=None, need_attn=False, ) # if model.args.fp16: # model.half() # hash model self.model_hash = hash_model(model_dir) if self.model_hash in MODELS: model_langs = MODELS[self.model_hash]['langs'] if lang not in model_langs: model_name = MODELS[self.model_hash]['name'] logger.warning( f'Language "{lang}" is unsupported for model "{model_name}"' ) logger.warning( f'Supported languages for {model_name}: {", ".join(model_langs)}' ) else: logger.warning('unrecognized model, so cannot check language')
def from_pretrained(cls, model_name_or_path, checkpoint_file='model.pt', data_name_or_path=None, **kwargs): """ Load a :class:`~fairseq.models.FairseqModel` from a pre-trained model file. Downloads and caches the pre-trained model file if needed. The base implementation returns a :class:`fairseq.hub_utils.Generator`, which can be used to generate translations or sample from language models. The underlying :class:`~fairseq.models.FairseqModel` can be accessed via the *generator.models* attribute. Other models may override this to implement custom PyTorch Hub APIs. Args: model_name_or_path (str): either the name of a pre-trained model to load or a path/URL to a pre-trained model state dict checkpoint_file (str, optional): colon-separated list of checkpoint files in the model archive to ensemble (default: 'model.pt') data_name_or_path (str, optional): point args.data to the archive at the given path/URL. Can start with '.' or './' to reuse the model archive path. """ from fairseq import checkpoint_utils, file_utils, hub_utils if hasattr(cls, 'hub_models'): archive_map = cls.hub_models() if model_name_or_path in archive_map: model_name_or_path = archive_map[model_name_or_path] if data_name_or_path is not None and data_name_or_path in archive_map: data_name_or_path = archive_map[data_name_or_path] model_path = file_utils.load_archive_file(model_name_or_path) # convenience hack for loading data and BPE codes from model archive if data_name_or_path is not None: if data_name_or_path.startswith('.'): kwargs['data'] = os.path.abspath(os.path.join(model_path, data_name_or_path)) else: kwargs['data'] = file_utils.load_archive_file(data_name_or_path) for file, arg in { 'code': 'bpe_codes', 'bpecodes': 'bpe_codes', 'sentencepiece.bpe.model': 'sentencepiece_vocab', }.items(): path = os.path.join(model_path, file) if os.path.exists(path): kwargs[arg] = path models, args, task = checkpoint_utils.load_model_ensemble_and_task( [os.path.join(model_path, cpt) for cpt in checkpoint_file.split(':')], arg_overrides=kwargs, ) print(args) return hub_utils.Generator(args, task, models)
def _load_model(self, model_path, target_dict): class BasicTask(object): # Class to get around creating a full task object. def __init__(self, target_dict): self.target_dictionary = target_dict data_folder = self.model_path.replace( self.model_path.split('/')[-1], '') args_overrides = { 'data': data_folder, } try: model, cfg, task = checkpoint_utils.load_model_ensemble_and_task( [model_path], arg_overrides=args_overrides) except MissingMandatoryValue as e: args_overrides['w2v_path'] = model_path model, cfg, task = checkpoint_utils.load_model_ensemble_and_task( [model_path], arg_overrides=args_overrides) return model
def download_and_load_checkpoint( self, checkpoint_filename: str, arg_overrides: Optional[Dict[str, str]] = None): path = self.download(self.base_url, self.root, checkpoint_filename) _arg_overrides = arg_overrides or {} _arg_overrides["data"] = self.root.as_posix() models, cfg, task = load_model_ensemble_and_task( [path.as_posix()], arg_overrides=_arg_overrides) if self.use_cuda: for model in models: model.cuda() generator = task.build_generator(models, cfg) return models, cfg, task, generator
def test_load_model_ensemble_and_task(self): with contextlib.redirect_stdout(StringIO()): with self._train_transformer(seed=123) as model1: with self._train_transformer(seed=456) as model2: ensemble, cfg, task = checkpoint_utils.load_model_ensemble_and_task( filenames=[model1, model2]) self.assertEqual(len(ensemble), 2) # after Transformer has been migrated to Hydra, this will probably # become cfg.common.seed self.assertEqual(ensemble[0].args.seed, 123) self.assertEqual(ensemble[1].args.seed, 456) # the task from the first model should be returned self.assertEqual(task.args.seed, 123)
def test_prune_state_dict(self): with contextlib.redirect_stdout(StringIO()): extra_args = [ "--encoder-layerdrop", "0.01", "--decoder-layerdrop", "0.01" ] with self._train_transformer(seed=1, extra_args=extra_args) as model: ensemble, cfg, task = checkpoint_utils.load_model_ensemble_and_task( filenames=[model], arg_overrides={ "encoder_layers_to_keep": "0,2", "decoder_layers_to_keep": "1", }, ) self.assertEqual(len(ensemble), 1) self.assertEqual(len(ensemble[0].encoder.layers), 2) self.assertEqual(len(ensemble[0].decoder.layers), 1)
def from_pretrained( model_name_or_path, checkpoint_file='model.pt', data_name_or_path='.', archive_map=None, **kwargs ): from fairseq import checkpoint_utils, file_utils if archive_map is not None: if model_name_or_path in archive_map: model_name_or_path = archive_map[model_name_or_path] if data_name_or_path is not None and data_name_or_path in archive_map: data_name_or_path = archive_map[data_name_or_path] model_path = file_utils.load_archive_file(model_name_or_path) # convenience hack for loading data and BPE codes from model archive if data_name_or_path.startswith('.'): kwargs['data'] = os.path.abspath(os.path.join(model_path, data_name_or_path)) else: kwargs['data'] = file_utils.load_archive_file(data_name_or_path) for file, arg in { 'code': 'bpe_codes', 'bpecodes': 'bpe_codes', 'sentencepiece.bpe.model': 'sentencepiece_vocab', }.items(): path = os.path.join(model_path, file) if os.path.exists(path): kwargs[arg] = path if 'user_dir' in kwargs: utils.import_user_module(argparse.Namespace(user_dir=kwargs['user_dir'])) models, args, task = checkpoint_utils.load_model_ensemble_and_task( [os.path.join(model_path, cpt) for cpt in checkpoint_file.split(':')], arg_overrides=kwargs, ) return { 'args': args, 'task': task, 'models': models, }
def main(): parser = argparse.ArgumentParser() # rescoring_task.RescoringTask.add_args(parser, add_def_opts=False, add_data_opts=False) WordTokenizer.add_args(parser) #LatsDataSet.add_args(parser) LatsOracleAlignDataSet.add_args(parser) parser.add_argument('--no_progress_bar', action='store_true', help="disable progress bar") parser.add_argument('--device', type=torch.device, default=torch.device('cpu'), help="PyTorch device") parser.add_argument('--format', type=str, default='pdf', help="Render file type. Default .pdf") parser.add_argument('--norm', action='store_true', default=False, help="Normalize lt probs") parser.add_argument("model", type=str, help="Saved model") parser.add_argument("out_dir", type=str, help='Out dir') parser.add_argument("utts", nargs="+", help="Utt_ids for draw") args = parser.parse_args() setup_logger() tokenizer = WordTokenizer.build_from_args(args) #dataset = LatsDataSet.build_from_args(args, tokenizer) dataset = LatsOracleAlignDataSet.build_from_args(args, tokenizer) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [args.model], arg_overrides=vars(args)) model = models[-1].to(args.device) logger.info(f"Drawing {args.utts}") draw_rescoring_lats(model, dataset, args.out_dir, args.utts, norm=args.norm, device=args.device, format=args.format, progress_bar=(not args.no_progress_bar))
def setup_ar(checkpoint, batch_size, tokens_per_target=None, fp16=False): models, args, task = checkpoint_utils.load_model_ensemble_and_task( [checkpoint]) if tokens_per_target is None: tokens_per_target = args.tokens_per_sample args_overrides = { 'cpu': False, 'no_beamable_mm': False, 'max_sentences': batch_size, 'fp16': fp16, 'tokens_per_target': tokens_per_target, 'sampling': True, 'beam': 1 } for arg_name, arg_val in args_overrides.items(): setattr(args, arg_name, arg_val) use_cuda = torch.cuda.is_available() and not args.cpu for model in models: model.make_generation_fast_( beamable_mm_beam_size=None if args.no_beamable_mm else args.max_sentences, need_attn=False, ) if args.fp16: model.half() if use_cuda: model.cuda() generator = task.build_generator(args) def generate_fn(n): bs = task.args.max_sentences results = [] for i in range(0, n, bs): real_bs = min(bs, n - i) sample = build_sample(task, bs=real_bs) tokens, _ = task.inference_step(generator, models, sample) results.append(tokens) return np.concatenate(results, axis=0) return generate_fn
def set_up_model(args, path, override_args=None): use_fp16 = args.fp16 use_cuda = torch.cuda.is_available() and not args.cpu if override_args is not None: overrides = vars(override_args) overrides.update(eval(getattr(override_args, 'model_overrides', '{}'))) else: overrides = None # Load ensemble print('| loading model(s) from {}'.format(path)) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [path], arg_overrides=overrides, ) model = models[0] # Move models to GPU for model in models: if use_fp16: model.half() if use_cuda: model.cuda() # Print args print(model_args) # Build criterion criterion = task.build_criterion(model_args) criterion.eval() if model_args.task == 'language_modeling': model_args.sampling = True model_args.sampling_topk = 20 model_args.max_len_b = 80 generator = task.build_generator(model_args) return model, task, criterion, generator, model_args
def load_models_and_task(args, use_cuda, task=None): # if `task` is not provided, it will be from the saved model args models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( args.path.split(':'), arg_overrides=eval(args.model_overrides), task=task, ) # Optimize ensemble for generation for model in models: model.make_generation_fast_( beamable_mm_beam_size=None if args.no_beamable_mm else args.beam, need_attn=args.print_alignment, ) if args.fp16: print('using fp16 for models') model.half() if use_cuda: print('using GPU for models') model.cuda() else: print('using CPU for models') # model = Model(models, task.target_dictionary) return models, model_args, task
type=torch.device, default=torch.device('cpu'), help="PyTorch device") parser.add_argument('--keep_tmp', action='store_true', help="keep tmp hyp") parser.add_argument('--acwt', type=float, default=1, help="Acoustic weight") parser.add_argument("model", type=str, help="Saved model") parser.add_argument("ref_fname", type=str, help="Reference text") args = parser.parse_args() setup_logger() tokenizer = WordTokenizer.build_from_args(args) dataset = LatsDataSet.build_from_args(args, tokenizer) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [args.model], arg_overrides=vars(args)) model = models[-1].to(args.device) ref_fname = args.ref_fname wer = compute_model_wer(model=model, dataset=dataset, ref_fname=ref_fname, device=args.device, btz=args.infer_btz, model_weight=args.model_weight, progress_bar=(not args.no_progress_bar), keep_tmp=args.keep_tmp, dataloader_nj=args.dataloader_nj, hyp_filter=args.hyp_filter, acwt=args.acwt, lmwt=args.lmwt, strategy=args.strategy)
def decode_rerank_scores(args): if args.max_tokens is None and args.batch_size is None: args.batch_size = 1 logger.info(args) use_cuda = torch.cuda.is_available() and not args.cpu # Load ensemble logger.info("loading model(s) from {}".format(args.path)) models, _model_args, task = checkpoint_utils.load_model_ensemble_and_task( [args.path], arg_overrides=eval(args.model_overrides), ) for model in models: if args.fp16: model.half() if use_cuda: model.cuda() # Initialize generator generator = task.build_generator(args) # Handle tokenization and BPE tokenizer = task.build_tokenizer(args) bpe = task.build_bpe(args) def encode_fn(x): if tokenizer is not None: x = tokenizer.encode(x) if bpe is not None: x = bpe.encode(x) return x max_positions = utils.resolve_max_positions( task.max_positions(), *[model.max_positions() for model in models]) src, hyp, mt_scores = parse_fairseq_gen(args.in_text, task) model_scores = {} logger.info("decode reranker score") for batch in make_batches(args, src, hyp, task, max_positions, encode_fn): src_tokens = batch.src_tokens src_lengths = batch.src_lengths if use_cuda: src_tokens = src_tokens.cuda() src_lengths = src_lengths.cuda() sample = { "net_input": { "src_tokens": src_tokens, "src_lengths": src_lengths }, } scores = task.inference_step(generator, models, sample) for id, sc in zip(batch.ids.tolist(), scores.tolist()): model_scores[id] = sc[0] model_scores = [model_scores[i] for i in range(len(model_scores))] return src, hyp, mt_scores, model_scores
def main(cfg: DictConfig, override_args=None): if isinstance(cfg, Namespace): cfg = convert_namespace_to_omegaconf(cfg) utils.import_user_module(cfg.common) assert ( cfg.dataset.max_tokens is not None or cfg.dataset.batch_size is not None ), "Must specify batch size either with --max-tokens or --batch-size" use_fp16 = cfg.common.fp16 use_cuda = torch.cuda.is_available() and not cfg.common.cpu if use_cuda: torch.cuda.set_device(cfg.distributed_training.device_id) if cfg.distributed_training.distributed_world_size > 1: data_parallel_world_size = distributed_utils.get_data_parallel_world_size( ) data_parallel_rank = distributed_utils.get_data_parallel_rank() else: data_parallel_world_size = 1 data_parallel_rank = 0 if override_args is not None: overrides = vars(override_args) overrides.update(eval(getattr(override_args, "model_overrides", "{}"))) else: overrides = None # Load ensemble logger.info("loading model(s) from {}".format(cfg.common_eval.path)) models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task( [cfg.common_eval.path], arg_overrides=overrides, suffix=cfg.checkpoint.checkpoint_suffix, ) model = models[0] # Move models to GPU for model in models: if use_fp16: model.half() if use_cuda: model.cuda() # Print args logger.info(saved_cfg) # Build criterion criterion = task.build_criterion(saved_cfg.criterion) criterion.eval() for subset in cfg.dataset.valid_subset.split(","): try: task.load_dataset(subset, combine=False, epoch=1, task_cfg=saved_cfg.task) dataset = task.dataset(subset) except KeyError: raise Exception("Cannot find dataset: " + subset) # Initialize data iterator itr = task.get_batch_iterator( dataset=dataset, max_tokens=cfg.dataset.max_tokens, max_sentences=cfg.dataset.batch_size, max_positions=utils.resolve_max_positions( task.max_positions(), *[m.max_positions() for m in models], ), ignore_invalid_inputs=cfg.dataset. skip_invalid_size_inputs_valid_test, required_batch_size_multiple=cfg.dataset. required_batch_size_multiple, seed=cfg.common.seed, num_shards=data_parallel_world_size, shard_id=data_parallel_rank, num_workers=cfg.dataset.num_workers, data_buffer_size=cfg.dataset.data_buffer_size, ).next_epoch_itr(shuffle=False) progress = progress_bar.progress_bar( itr, log_format=cfg.common.log_format, log_interval=cfg.common.log_interval, prefix=f"valid on '{subset}' subset", default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"), ) log_outputs = [] for i, sample in enumerate(progress): sample = utils.move_to_cuda(sample) if use_cuda else sample _loss, _sample_size, log_output = task.valid_step( sample, model, criterion) progress.log(log_output, step=i) log_outputs.append(log_output) if data_parallel_world_size > 1: log_outputs = distributed_utils.all_gather_list( log_outputs, max_size=cfg.common.all_gather_list_size, group=distributed_utils.get_data_parallel_group(), ) log_outputs = list(chain.from_iterable(log_outputs)) with metrics.aggregate() as agg: task.reduce_metrics(log_outputs, criterion) log_output = agg.get_smoothed_values() progress.print(log_output, tag=subset, step=i)
def main(args, override_args=None): utils.import_user_module(args) use_fp16 = args.fp16 use_cuda = torch.cuda.is_available() and not args.cpu if override_args is not None: overrides = vars(override_args) overrides.update(eval(getattr(override_args, 'model_overrides', '{}'))) else: overrides = None # Load ensemble print('| loading model(s) from {}'.format(args.path)) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [args.path], arg_overrides=overrides, ) model = models[0] # Move models to GPU for model in models: if use_fp16: model.half() if use_cuda: model.cuda() # Print args print(model_args) # Build criterion criterion = task.build_criterion(model_args) criterion.eval() # Load valid dataset (we load training data below, based on the latest checkpoint) for subset in args.valid_subset.split(','): try: task.load_dataset(subset, combine=False, epoch=0) dataset = task.dataset(subset) except KeyError: raise Exception('Cannot find dataset: ' + subset) # Initialize data iterator itr = task.get_batch_iterator( dataset=dataset, max_tokens=args.max_tokens, max_sentences=args.max_sentences, max_positions=utils.resolve_max_positions( task.max_positions(), *[m.max_positions() for m in models], ), ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test, required_batch_size_multiple=args.required_batch_size_multiple, seed=args.seed, num_workers=args.num_workers, ).next_epoch_itr(shuffle=False) progress = progress_bar.build_progress_bar( args, itr, prefix='valid on \'{}\' subset'.format(subset), no_progress_bar='simple') log_outputs = [] for i, sample in enumerate(progress): sample = utils.move_to_cuda(sample) if use_cuda else sample _loss, _sample_size, log_output = task.valid_step( sample, model, criterion) progress.log(log_output, step=i) log_outputs.append(log_output) log_output = task.aggregate_logging_outputs(log_outputs, criterion) progress.print(log_output, tag=subset, step=i)
def setup_ar(checkpoint, batch_size, tokens_per_target=None, fp16=False): models, args, task = checkpoint_utils.load_model_ensemble_and_task([checkpoint]) if tokens_per_target is None:
def main(rank, world_size, args): if world_size > 1: torch.distributed.init_process_group(backend="gloo", init_method="env://", world_size=world_size, rank=rank) torch.cuda.set_device(rank) raw_args = args args = convert_namespace_to_omegaconf(args) if args.common.seed is not None: random.seed(args.common.seed) np.random.seed(args.common.seed) utils.set_torch_seed(args.common.seed) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [raw_args.path], arg_overrides={"data": args.task.data}) tgt_dict = task.target_dictionary for model in models: model.prepare_for_inference_(args) model.cuda().eval() if raw_args.fp16: model = model.half() model = models[0] config = ExpressiveCodeDataConfig(args.task.data) dataset = CodeDataset( manifest=config.manifests[raw_args.subset], dictionary=task.source_dictionary, dur_dictionary=task.source_duration_dictionary, f0_dictionary=task.source_f0_dictionary, config=config, discrete_dur=task.cfg.discrete_duration, discrete_f0=task.cfg.discrete_f0, log_f0=task.cfg.log_f0, normalize_f0_mean=task.cfg.normalize_f0_mean, normalize_f0_std=task.cfg.normalize_f0_std, interpolate_f0=task.cfg.interpolate_f0, shifts=task.cfg.stream_shifts, return_filename=True, strip_filename=False, ) tgt_dict = task.target_dictionary shifts = dataset.shifts.dur, dataset.shifts.f0 max_shift = max(shifts) fname = raw_args.output if world_size > 1: fname += f"_{rank}" output_file = open(fname, "w") if raw_args.filter_names: dataset = FilterNamesDataset(dataset, raw_args.filter_names) dataset = InferenceDataset(dataset, raw_args.prefix_length, filter_short=True) print(f"Dataset size {len(dataset)}") sampler = (None if world_size == 1 else DistributedSampler( dataset, num_replicas=world_size, rank=rank, shuffle=False)) dataloader = DataLoader( dataset, batch_size=1, shuffle=False, collate_fn=dataset.collater, sampler=sampler, ) Ts = raw_args.T_token, raw_args.T_duration, raw_args.T_f0 decoder = TemperatureDecoder(Ts, discrete_dur=task.cfg.discrete_duration, discrete_f0=task.cfg.discrete_f0) dataset_size = len(dataset) f0_decoder = None if raw_args.f0_discretization_bounds: assert task.cfg.discrete_f0 f0_decoder = Naive_F0_Decoder(raw_args.f0_discretization_bounds).cuda() pbar = (tqdm.tqdm( total=dataset_size if raw_args.max_samples is None else min( raw_args.max_samples, dataset_size)) if world_size == 1 else None) samples_produced = 0 for batch in dataloader: if (raw_args.max_samples is not None and samples_produced >= raw_args.max_samples): break prefix = batch["prefix"][0] batch = explode_batch(batch, raw_args.batch_explosion_rate) batch = move_to_cuda(batch) if not raw_args.short_curcuit: produced_tokens, produced_durations, produced_f0, _ = do_sampling( models[0], batch, tgt_dict.eos(), decoder, autoregressive_steps=raw_args.max_length - prefix + max_shift, teacher_force_tokens=raw_args.teacher_force_tokens, match_duration=raw_args.match_duration, teacher_force_duration=raw_args.teacher_force_duration, teacher_force_f0=raw_args.teacher_force_f0, ) # stip entries corresponding to <s> produced_tokens = produced_tokens[:, 1:] produced_durations = produced_durations[:, 1:] produced_f0 = produced_f0[:, 1:] else: max_length = raw_args.max_length + max_shift produced_tokens, produced_durations, produced_f0 = ( batch["target"][:, :max_length], batch["dur_target"][:, :max_length], batch["f0_target"][:, :max_length], ) if f0_decoder is not None: produced_f0 = f0_decoder(produced_f0) produced_tokens, produced_durations, produced_f0 = ( produced_tokens.cpu().tolist(), produced_durations.cpu().tolist(), produced_f0.cpu().tolist(), ) bsz = batch["target"].size(0) assert bsz == raw_args.batch_explosion_rate for i in range(bsz): if (raw_args.max_samples is not None and samples_produced >= raw_args.max_samples): break produced_tokens_i = produced_tokens[i] produced_durations_i = produced_durations[i] produced_f0_i = produced_f0[i] ( produced_tokens_i, produced_durations_i, produced_f0_i, ) = realign_shifted_streams(produced_tokens_i, produced_durations_i, produced_f0_i, shifts) produced_tokens_i, produced_durations_i, produced_f0_i = maybe_cut_eos( produced_tokens_i, produced_durations_i, produced_f0_i, tgt_dict.eos()) produced_tokens_i, produced_durations_i, produced_f0_i = maybe_filter_pad( produced_tokens_i, produced_durations_i, produced_f0_i, tgt_dict.pad()) if raw_args.match_duration: # NB: here we cheat a bit and use that padding has duration 0 # so no need to re-align and remove padding dur_target_i = batch["dur_target"][i, :].sum().item() produced_tokens_i, produced_durations_i, produced_f0_i = match_duration( produced_tokens_i, produced_durations_i, produced_f0_i, dur_target_i) if raw_args.cut_prompt: produced_tokens_i, produced_durations_i, produced_f0_i = ( produced_tokens_i[prefix:], produced_durations_i[prefix:], produced_f0_i[prefix:], ) prompt_fname = batch["filename"][0] fname = str( pathlib.Path(prompt_fname).with_suffix("")) + f"__{i}.wav" token_stream = unroll_duration(produced_tokens_i, produced_durations_i) f0_stream = unroll_duration(produced_f0_i, produced_durations_i) output_line = json.dumps({ "audio": fname, "prompt": prompt_fname, raw_args.code_type: " ".join(map(str, token_stream)), "duration": round( sum(produced_durations_i) * CODETYPE_TO_FRAMETIME[raw_args.code_type], 3, ), "raw_duration": produced_durations_i, "raw_f0": produced_f0_i, "f0": [round(f0, 3) for f0 in f0_stream], }) print(output_line, file=output_file) if pbar: pbar.update(1) samples_produced += 1 if raw_args.debug: break output_file.close() if world_size > 1: # important that everything is flushed before aggregating torch.distributed.barrier() if world_size > 1 and rank == 0: with open(raw_args.output, "w") as fout: for i in range(world_size): f = raw_args.output + f"_{i}" with open(f, "r") as fin: fout.write(fin.read()) os.remove(f)
def main(args, task=None, model_state=None): check_args(args) if args.max_tokens is None and args.batch_size is None: args.max_tokens = 4000000 logger.info(args) use_cuda = torch.cuda.is_available() and not args.cpu logger.info("| decoding with criterion {}".format(args.criterion)) task = tasks.setup_task(args) # Load ensemble if args.load_emissions: models, criterions = [], [] task.load_dataset(args.gen_subset) else: logger.info("| loading model(s) from {}".format(args.path)) models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task( utils.split_paths(args.path), arg_overrides=ast.literal_eval(args.model_overrides), task=task, suffix=args.checkpoint_suffix, strict=(args.checkpoint_shard_count == 1), num_shards=args.checkpoint_shard_count, state=model_state, ) optimize_models(args, use_cuda, models) task.load_dataset(args.gen_subset, task_cfg=saved_cfg.task) # Set dictionary tgt_dict = task.target_dictionary logger.info( "| {} {} {} examples".format( args.data, args.gen_subset, len(task.dataset(args.gen_subset)) ) ) # hack to pass transitions to W2lDecoder if args.criterion == "asg_loss": raise NotImplementedError("asg_loss is currently not supported") # trans = criterions[0].asg.trans.data # args.asg_transitions = torch.flatten(trans).tolist() # Load dataset (possibly sharded) itr = get_dataset_itr(args, task, models) # Initialize generator gen_timer = StopwatchMeter() def build_generator(args): w2l_decoder = getattr(args, "w2l_decoder", None) if w2l_decoder == "viterbi": from examples.speech_recognition.w2l_decoder import W2lViterbiDecoder return W2lViterbiDecoder(args, task.target_dictionary) elif w2l_decoder == "kenlm": from examples.speech_recognition.w2l_decoder import W2lKenLMDecoder return W2lKenLMDecoder(args, task.target_dictionary) elif w2l_decoder == "fairseqlm": from examples.speech_recognition.w2l_decoder import W2lFairseqLMDecoder return W2lFairseqLMDecoder(args, task.target_dictionary) else: print( "only flashlight decoders with (viterbi, kenlm, fairseqlm) options are supported at the moment" ) # please do not touch this unless you test both generate.py and infer.py with audio_pretraining task generator = build_generator(args) if args.load_emissions: generator = ExistingEmissionsDecoder( generator, np.load(args.load_emissions, allow_pickle=True) ) logger.info("loaded emissions from " + args.load_emissions) num_sentences = 0 if args.results_path is not None and not os.path.exists(args.results_path): os.makedirs(args.results_path) max_source_pos = ( utils.resolve_max_positions( task.max_positions(), *[model.max_positions() for model in models] ), ) if max_source_pos is not None: max_source_pos = max_source_pos[0] if max_source_pos is not None: max_source_pos = max_source_pos[0] - 1 if args.dump_emissions: emissions = {} if args.dump_features: features = {} models[0].bert.proj = None else: res_files = prepare_result_files(args) errs_t = 0 lengths_t = 0 with progress_bar.build_progress_bar(args, itr) as t: wps_meter = TimeMeter() for sample in t: sample = utils.move_to_cuda(sample) if use_cuda else sample if "net_input" not in sample: continue prefix_tokens = None if args.prefix_size > 0: prefix_tokens = sample["target"][:, : args.prefix_size] gen_timer.start() if args.dump_emissions: with torch.no_grad(): encoder_out = models[0](**sample["net_input"]) emm = models[0].get_normalized_probs(encoder_out, log_probs=True) emm = emm.transpose(0, 1).cpu().numpy() for i, id in enumerate(sample["id"]): emissions[id.item()] = emm[i] continue elif args.dump_features: with torch.no_grad(): encoder_out = models[0](**sample["net_input"]) feat = encoder_out["encoder_out"].transpose(0, 1).cpu().numpy() for i, id in enumerate(sample["id"]): padding = ( encoder_out["encoder_padding_mask"][i].cpu().numpy() if encoder_out["encoder_padding_mask"] is not None else None ) features[id.item()] = (feat[i], padding) continue hypos = task.inference_step(generator, models, sample, prefix_tokens) num_generated_tokens = sum(len(h[0]["tokens"]) for h in hypos) gen_timer.stop(num_generated_tokens) for i, sample_id in enumerate(sample["id"].tolist()): speaker = None # id = task.dataset(args.gen_subset).ids[int(sample_id)] id = sample_id toks = ( sample["target"][i, :] if "target_label" not in sample else sample["target_label"][i, :] ) target_tokens = utils.strip_pad(toks, tgt_dict.pad()).int().cpu() # Process top predictions errs, length = process_predictions( args, hypos[i], None, tgt_dict, target_tokens, res_files, speaker, id, ) errs_t += errs lengths_t += length wps_meter.update(num_generated_tokens) t.log({"wps": round(wps_meter.avg)}) num_sentences += ( sample["nsentences"] if "nsentences" in sample else sample["id"].numel() ) wer = None if args.dump_emissions: emm_arr = [] for i in range(len(emissions)): emm_arr.append(emissions[i]) np.save(args.dump_emissions, emm_arr) logger.info(f"saved {len(emissions)} emissions to {args.dump_emissions}") elif args.dump_features: feat_arr = [] for i in range(len(features)): feat_arr.append(features[i]) np.save(args.dump_features, feat_arr) logger.info(f"saved {len(features)} emissions to {args.dump_features}") else: if lengths_t > 0: wer = errs_t * 100.0 / lengths_t logger.info(f"WER: {wer}") logger.info( "| Processed {} sentences ({} tokens) in {:.1f}s ({:.2f}" "sentences/s, {:.2f} tokens/s)".format( num_sentences, gen_timer.n, gen_timer.sum, num_sentences / gen_timer.sum, 1.0 / gen_timer.avg, ) ) logger.info("| Generate {} with beam={}".format(args.gen_subset, args.beam)) return task, wer
def main(args, override_args=None): utils.import_user_module(args) assert args.max_tokens is not None or args.max_sentences is not None, \ 'Must specify batch size either with --max-tokens or --max-sentences' use_fp16 = args.fp16 use_cuda = torch.cuda.is_available() and not args.cpu if use_cuda: torch.cuda.set_device(args.device_id) if override_args is not None: overrides = vars(override_args) overrides.update(eval(getattr(override_args, 'model_overrides', '{}'))) else: overrides = None # Load ensemble logger.info('loading model(s) from {}'.format(args.path)) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [args.path], arg_overrides=overrides, suffix=getattr(args, "checkpoint_suffix", ""), ) model = models[0] # Move models to GPU for model in models: if use_fp16: model.half() if use_cuda: model.cuda() # Print args logger.info(model_args) # Build criterion criterion = task.build_criterion(model_args) criterion.eval() for subset in args.valid_subset.split(','): try: task.load_dataset(subset, combine=False, epoch=1) dataset = task.dataset(subset) except KeyError: raise Exception('Cannot find dataset: ' + subset) # Initialize data iterator itr = task.get_batch_iterator( dataset=dataset, max_tokens=args.max_tokens, max_sentences=args.max_sentences, max_positions=utils.resolve_max_positions( task.max_positions(), *[m.max_positions() for m in models], ), ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test, required_batch_size_multiple=args.required_batch_size_multiple, seed=args.seed, num_shards=args.distributed_world_size, shard_id=args.distributed_rank, num_workers=args.num_workers, data_buffer_size=args.data_buffer_size, ).next_epoch_itr(shuffle=False) progress = progress_bar.progress_bar( itr, log_format=args.log_format, log_interval=args.log_interval, prefix=f"valid on '{subset}' subset", default_log_format=('tqdm' if not args.no_progress_bar else 'simple'), ) log_outputs = [] for i, sample in enumerate(progress): sample = utils.move_to_cuda(sample) if use_cuda else sample _loss, _sample_size, log_output = task.valid_step( sample, model, criterion) progress.log(log_output, step=i) log_outputs.append(log_output) if args.distributed_world_size > 1: log_outputs = distributed_utils.all_gather_list( log_outputs, max_size=getattr(args, 'all_gather_list_size', 16384), ) log_outputs = list(chain.from_iterable(log_outputs)) with metrics.aggregate() as agg: task.reduce_metrics(log_outputs, criterion) log_output = agg.get_smoothed_values() progress.print(log_output, tag=subset, step=i)
def main(cfg: DictConfig, override_args=None, **unused_kwargs): if isinstance(cfg, Namespace): cfg = convert_namespace_to_omegaconf(cfg) utils.import_user_module(cfg.common) use_fp16 = cfg.common.fp16 use_cuda = torch.cuda.is_available() and not cfg.common.cpu if use_cuda: torch.cuda.set_device(cfg.distributed_training.device_id) if override_args is not None: overrides = vars(override_args) overrides.update(eval(getattr(override_args, "model_overrides", "{}"))) else: overrides = None logger.info(cfg) # Load ensemble logger.info("loading model(s) from {}".format(cfg.common_eval.path)) # reduce tokens per sample by the required context window size cfg.task.tokens_per_sample -= cfg.eval_lm.context_window models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [cfg.common_eval.path], arg_overrides=overrides, suffix=cfg.checkpoint.checkpoint_suffix, strict=(cfg.checkpoint.checkpoint_shard_count == 1), num_shards=cfg.checkpoint.checkpoint_shard_count, ) # Load dataset splits gen_subset = cfg.dataset.gen_subset task.load_dataset(gen_subset) dataset = task.dataset(gen_subset) if cfg.eval_lm.context_window > 0: dataset = LMContextWindowDataset( dataset=dataset, tokens_per_sample=cfg.task.tokens_per_sample, context_window=cfg.eval_lm.context_window, pad_idx=task.source_dictionary.pad(), ) logger.info("{} {} {} examples".format(cfg.task.data, gen_subset, len(dataset))) # Optimize ensemble for generation and set the source and dest dicts on the model (required by scorer) for model in models: if use_fp16: model.half() if use_cuda and not cfg.distributed_training.pipeline_model_parallel: model.cuda() model.prepare_for_inference_(cfg) assert len(models) > 0 logger.info( "num. model params: {}".format(sum(p.numel() for p in models[0].parameters())) ) itr = task.get_batch_iterator( dataset=dataset, max_tokens=cfg.dataset.max_tokens or 36000, max_sentences=cfg.dataset.batch_size, max_positions=utils.resolve_max_positions( *[model.max_positions() for model in models] ), ignore_invalid_inputs=True, num_shards=max( cfg.dataset.num_shards, cfg.distributed_training.distributed_world_size, ), shard_id=max( cfg.dataset.shard_id, cfg.distributed_training.distributed_rank, ), num_workers=cfg.dataset.num_workers, data_buffer_size=cfg.dataset.data_buffer_size, ).next_epoch_itr(shuffle=False) progress = progress_bar.progress_bar( itr, log_format=cfg.common.log_format, log_interval=cfg.common.log_interval, default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"), ) gen_timer = StopwatchMeter() scorer = SequenceScorer(task.target_dictionary, cfg.eval_lm.softmax_batch) score_sum = 0.0 count = 0 if cfg.common_eval.remove_bpe is not None: if cfg.common_eval.remove_bpe == "sentencepiece": raise NotImplementedError else: bpe_cont = cfg.common_eval.remove_bpe.rstrip() bpe_toks = { i for i in range(len(task.source_dictionary)) if task.source_dictionary[i].endswith(bpe_cont) } bpe_len = len(bpe_cont) else: bpe_toks = None bpe_len = 0 word_stats = dict() wps_meter = TimeMeter() for sample in progress: if "net_input" not in sample: continue sample = utils.move_to_cuda(sample) if use_cuda else sample gen_timer.start() hypos = scorer.generate(models, sample) gen_timer.stop(sample["ntokens"]) for i, hypos_i in enumerate(hypos): hypo = hypos_i[0] sample_id = sample["id"][i] tokens = hypo["tokens"] tgt_len = tokens.numel() pos_scores = hypo["positional_scores"].float() if cfg.task.add_bos_token: assert hypo["tokens"][0].item() == task.target_dictionary.bos() tokens = tokens[1:] pos_scores = pos_scores[1:] skipped_toks = 0 if bpe_toks is not None: for i in range(tgt_len - 1): if tokens[i].item() in bpe_toks: skipped_toks += 1 pos_scores[i + 1] += pos_scores[i] pos_scores[i] = 0 inf_scores = pos_scores.eq(float("inf")) | pos_scores.eq(float("-inf")) if inf_scores.any(): logger.info( "skipping tokens with inf scores:", task.target_dictionary.string(tokens[inf_scores.nonzero()]), ) pos_scores = pos_scores[(~inf_scores).nonzero()] score_sum += pos_scores.sum().cpu() count += pos_scores.numel() - skipped_toks if cfg.eval_lm.output_word_probs or cfg.eval_lm.output_word_stats: w = "" word_prob = [] is_bpe = False for i in range(len(tokens)): w_ind = tokens[i].item() w += task.source_dictionary[w_ind] if bpe_toks is not None and w_ind in bpe_toks: w = w[:-bpe_len] is_bpe = True else: word_prob.append((w, pos_scores[i].item())) next_prob = None ind = i + 1 while ind < len(tokens): if pos_scores[ind].item() != 0: next_prob = pos_scores[ind] break ind += 1 word_stats.setdefault(w, WordStat(w, is_bpe)).add( pos_scores[i].item(), next_prob ) is_bpe = False w = "" if cfg.eval_lm.output_word_probs: logger.info( str(int(sample_id)) + " " + ( "\t".join( "{} [{:2f}]".format(x[0], x[1]) for x in word_prob ) ) ) wps_meter.update(sample["ntokens"]) progress.log({"wps": round(wps_meter.avg)}) avg_nll_loss = -score_sum / count / math.log(2) # convert to base 2 logger.info( "Evaluated {} tokens in {:.1f}s ({:.2f} tokens/s)".format( gen_timer.n, gen_timer.sum, 1.0 / gen_timer.avg ) ) logger.info( "Loss (base 2): {:.4f}, Perplexity: {:.2f}".format( avg_nll_loss, 2 ** avg_nll_loss ) ) if cfg.eval_lm.output_word_stats: for ws in sorted(word_stats.values(), key=lambda x: x.count, reverse=True): logger.info(ws)
def main(rank, world_size, args): start = time.time() if world_size > 1: torch.distributed.init_process_group(backend="gloo", init_method="env://", world_size=world_size, rank=rank) torch.cuda.set_device(rank % torch.cuda.device_count()) raw_args = args args = convert_namespace_to_omegaconf(args) if args.common.seed is not None: np.random.seed(args.common.seed) utils.set_torch_seed(args.common.seed) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [raw_args.path], arg_overrides={"data": args.task.data}) tgt_dict = task.target_dictionary for model in models: model.prepare_for_inference_(args) model.cuda().eval() if raw_args.fp16: model = model.half() model = models[0] config = ExpressiveCodeDataConfig(args.task.data) dataset = CodeDataset( manifest=config.manifests[raw_args.eval_subset], dictionary=task.source_dictionary, dur_dictionary=task.source_duration_dictionary, f0_dictionary=task.source_f0_dictionary, config=config, discrete_dur=task.cfg.discrete_duration, discrete_f0=task.cfg.discrete_f0, log_f0=task.cfg.log_f0, normalize_f0_mean=task.cfg.normalize_f0_mean, normalize_f0_std=task.cfg.normalize_f0_std, interpolate_f0=task.cfg.interpolate_f0, shifts=task.cfg.stream_shifts, return_filename=True, strip_filename=False, return_continuous_f0=raw_args.dequantize_prosody, ) if raw_args.filter_names: dataset = FilterNamesDataset(dataset, raw_args.filter_names) criterion = task.build_criterion(model_args.criterion) name2metric = { "continuation": continuation, "teacher_force_everything": teacher_force_everything, "correlation": correlation, } name2keys = { "continuation": ( "Token BLEU3", "Duration NLL", "Duration MAE", "F0 NLL", "F0 MAE", "F0 sum", "F0 sum_sq", "Dur sum", "Dur sum_sq", ), "teacher_force_everything": ("token_loss", "duration_loss", "f0_loss"), "correlation": ("Duration corr", "F0 corr"), } metric_name = raw_args.metric metric = name2metric[metric_name] results = metric(raw_args, dataset, model, criterion, tgt_dict, rank, world_size) values = None if metric_name not in [ "correlation", ]: values, normalizers = results values = maybe_aggregate_normalize(values, normalizers, world_size) elif metric_name == "correlation": values = maybe_aggregate_correlations(results, world_size) else: assert False assert values is not None summary = dict(zip(name2keys[raw_args.metric], values.tolist())) if metric_name == "continuation": summary["F0 Std"] = np.sqrt(-summary["F0 sum"]**2 + summary["F0 sum_sq"]) summary["Dur Std"] = np.sqrt(-summary["Dur sum"]**2 + summary["Dur sum_sq"]) del summary["F0 sum"] del summary["F0 sum_sq"] del summary["Dur sum"] del summary["Dur sum_sq"] summary["metric"] = metric_name if rank == 0: print(summary) if raw_args.wandb: wandb_results(summary, raw_args) print("# finished in ", time.time() - start, "seconds")
def main(cfg: DictConfig, **unused_kwargs): if isinstance(cfg, Namespace): cfg = convert_namespace_to_omegaconf(cfg) utils.import_user_module(cfg.common) logger.info(cfg) if cfg.eval_lm.context_window > 0: # reduce tokens per sample by the required context window size cfg.task.tokens_per_sample -= cfg.eval_lm.context_window # Initialize the task using the current *cfg* task = tasks.setup_task(cfg.task) # Load ensemble logger.info("loading model(s) from {}".format(cfg.common_eval.path)) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [cfg.common_eval.path], arg_overrides=eval(cfg.common_eval.model_overrides), suffix=cfg.checkpoint.checkpoint_suffix, strict=(cfg.checkpoint.checkpoint_shard_count == 1), num_shards=cfg.checkpoint.checkpoint_shard_count, task=task, ) use_fp16 = cfg.common.fp16 use_cuda = torch.cuda.is_available() and not cfg.common.cpu if use_cuda: torch.cuda.set_device(cfg.distributed_training.device_id) # Optimize ensemble for generation and set the source and dest dicts on the model # (required by scorer) for model in models: if use_fp16: model.half() if use_cuda and not cfg.distributed_training.pipeline_model_parallel: model.cuda() model.prepare_for_inference_(cfg) assert len(models) > 0 logger.info( "num. model params: {:,}".format(sum(p.numel() for p in models[0].parameters())) ) # Load dataset splits task.load_dataset(cfg.dataset.gen_subset) dataset = task.dataset(cfg.dataset.gen_subset) logger.info( "{} {} {:,} examples".format( cfg.task.data, cfg.dataset.gen_subset, len(dataset) ) ) itr = task.eval_lm_dataloader( dataset=dataset, max_tokens=cfg.dataset.max_tokens or 36000, batch_size=cfg.dataset.batch_size, max_positions=utils.resolve_max_positions( *[model.max_positions() for model in models] ), num_shards=max( cfg.dataset.num_shards, cfg.distributed_training.distributed_world_size, ), shard_id=max( cfg.dataset.shard_id, cfg.distributed_training.distributed_rank, ), num_workers=cfg.dataset.num_workers, data_buffer_size=cfg.dataset.data_buffer_size, context_window=cfg.eval_lm.context_window, ) itr = progress_bar.progress_bar( itr, log_format=cfg.common.log_format, log_interval=cfg.common.log_interval, default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"), ) results = eval_lm( models=models, source_dictionary=task.source_dictionary, batch_iterator=itr, post_process=cfg.common_eval.post_process, output_word_probs=cfg.eval_lm.output_word_probs, output_word_stats=cfg.eval_lm.output_word_stats, target_dictionary=task.target_dictionary, softmax_batch=cfg.eval_lm.softmax_batch, remove_bos_token=getattr(cfg.task, "add_bos_token", False), ) logger.info( "Loss (base 2): {:.4f}, Perplexity: {:.2f}".format( results["loss"], results["perplexity"] ) ) return results
def main(args, override_args, init_distributed=False): utils.import_user_module(args) assert args.max_tokens is not None or args.max_sentences is not None, \ 'Must specify batch size either with --max-tokens or --max-sentences' # Initialize CUDA and distributed training if torch.cuda.is_available() and not args.cpu: torch.cuda.set_device(args.device_id) np.random.seed(args.seed) torch.manual_seed(args.seed) if init_distributed: args.distributed_rank = distributed_utils.distributed_init(args) # Print args print(args) use_fp16 = args.fp16 use_cuda = torch.cuda.is_available() and not args.cpu if override_args is not None: overrides = vars(override_args) overrides.update(eval(getattr(override_args, 'model_overrides', '{}'))) else: overrides = None # Load ensemble print('| loading model(s) from {}'.format(args.path)) models, model_args, task = checkpoint_utils.load_model_ensemble_and_task( [args.path], arg_overrides=overrides, ) model = models[0] # Move models to GPU for model in models: if use_fp16: model.half() if use_cuda: model.cuda() # Print args print(model_args) # Build criterion criterion = task.build_criterion(model_args) criterion.eval() # Load valid dataset (we load training data below, based on the latest checkpoint) for valid_sub_split in args.valid_subset.split(','): task.load_dataset(valid_sub_split, combine=False, epoch=0) # Build trainer trainer = validator(args, task, model, criterion) print('| training on {} GPUs'.format(args.distributed_world_size)) print('| max tokens per GPU = {} and max sentences per GPU = {}'.format( args.max_tokens, args.max_sentences, )) # Load the latest checkpoint if one is available and restore the # corresponding train iterator valid_subsets = args.valid_subset.split(',') valid_losses = validate(args, trainer, task, valid_subsets)