class MTBottleneckConfig(NemoConfig): name: Optional[str] = 'MTBottleneck' do_training: bool = True do_testing: bool = False model: MTBottleneckModelConfig = MTBottleneckModelConfig() trainer: Optional[TrainerConfig] = TrainerConfig() exp_manager: Optional[ExpManagerConfig] = ExpManagerConfig(name='MTBottleneck', files_to_copy=[])
class MTEncDecConfig(NemoConfig): model: AAYNBaseConfig = AAYNBaseConfig() trainer: Optional[TrainerConfig] = TrainerConfig() exp_manager: Optional[ExpManagerConfig] = ExpManagerConfig( name='MTEncDec', files_to_copy=[])
class MTEncDecConfig(NemoConfig): name: Optional[str] = 'MTEncDec' do_training: bool = True model: MTEncDecModelConfig = MTEncDecModelConfig() trainer: Optional[TrainerConfig] = TrainerConfig() exp_manager: Optional[ExpManagerConfig] = ExpManagerConfig(name='MTEncDec', files_to_copy=[])
def nemo_export(argv): args = get_args(argv) loglevel = logging.INFO # assuming loglevel is bound to the string value obtained from the # command line argument. Convert to upper case to allow the user to # specify --log=DEBUG or --log=debug if args.verbose is not None: numeric_level = getattr(logging, args.verbose.upper(), None) if not isinstance(numeric_level, int): raise ValueError('Invalid log level: %s' % numeric_level) loglevel = numeric_level logger = logging.getLogger(__name__) if logger.handlers: for handler in logger.handlers: logger.removeHandler(handler) logging.basicConfig(level=loglevel, format='%(asctime)s [%(levelname)s] %(message)s') logging.info("Logging level set to {}".format(loglevel)) """Convert a .nemo saved model into .riva Riva input format.""" nemo_in = args.source out = args.out # Create a PL trainer object which is required for restoring Megatron models cfg_trainer = TrainerConfig( gpus=1, accelerator="ddp", num_nodes=1, # Need to set the following two to False as ExpManager will take care of them differently. logger=False, checkpoint_callback=False, ) trainer = Trainer(cfg_trainer) logging.info("Restoring NeMo model from '{}'".format(nemo_in)) try: with torch.inference_mode(): # Restore instance from .nemo file using generic model restore_from model = ModelPT.restore_from(restore_path=nemo_in, trainer=trainer) except Exception as e: logging.error( "Failed to restore model from NeMo file : {}. Please make sure you have the latest NeMo package installed with [all] dependencies.".format( nemo_in ) ) raise e logging.info("Model {} restored from '{}'".format(model.cfg.target, nemo_in)) if not isinstance(model, Exportable): logging.error("Your NeMo model class ({}) is not Exportable.".format(model.cfg.target)) sys.exit(1) typecheck.set_typecheck_enabled(enabled=False) try: # # Add custom export parameters here # in_args = {} if args.max_batch is not None: in_args["max_batch"] = args.max_batch if args.max_dim is not None: in_args["max_dim"] = args.max_dim autocast = nullcontext model = model.to(device=args.device) model.eval() with torch.inference_mode(): input_example = model.input_module.input_example(**in_args) if args.autocast: autocast = torch.cuda.amp.autocast with autocast(), torch.inference_mode(): logging.info(f"Getting output example") input_list, input_dict = parse_input_example(input_example) output_example = forward_method(model)(*input_list, **input_dict) logging.info(f"Exporting model with autocast={args.autocast}") input_names = model.input_names output_names = model.output_names _, descriptions = model.export( out, check_trace=False, input_example=input_example, onnx_opset_version=args.onnx_opset, verbose=args.verbose, ) except Exception as e: logging.error( "Export failed. Please make sure your NeMo model class ({}) has working export() and that you have the latest NeMo package installed with [all] dependencies.".format( model.cfg.target ) ) raise e logging.info("Successfully exported to {}".format(out)) del model if args.runtime_check: verify_runtime(out, input_list, input_dict, input_names, output_names, output_example)