def get_logits( snapshot_path: str, use_cuda_if_available: bool, output_path: Optional[str] = None, test_path: Optional[str] = None, field_names: Optional[List[str]] = None, ): _set_cuda(use_cuda_if_available) task, train_config = load(snapshot_path) print(f"Successfully loaded model from {snapshot_path}") if isinstance(task, NewTask): task.model.eval() data_source = _get_data_source( test_path, getattr(train_config.task.data, "source", None), field_names, task, ) task.data.batcher = Batcher() task.data.sort_key = None batches = task.data.batches(Stage.TEST, data_source=data_source) results = [] for (_, tensor_dict) in batches: model_inputs = task.model.arrange_model_inputs(tensor_dict) model_outputs = task.model(*model_inputs) MetricReporter.aggregate_data(results, model_outputs) with open(output_path, "w", encoding="utf-8") as fout: for row in results: fout.write(f"{row}\n")
def get_logits( snapshot_path: str, use_cuda_if_available: bool, output_path: Optional[str] = None, test_path: Optional[str] = None, field_names: Optional[List[str]] = None, dump_raw_input: bool = False, batch_size: int = 16, ndigits_precision: int = 0, output_columns: Optional[List[int]] = None, use_gzip: bool = False, device_id: int = 0, ): _set_cuda(use_cuda_if_available, device_id) task, train_config, _training_state = load(snapshot_path) print(f"Successfully loaded model from {snapshot_path}") print(f"Model on GPU? {next(task.model.parameters()).is_cuda}") if isinstance(task, NewTask): task.model.eval() data_source = _get_data_source(test_path, train_config.task.data, field_names, task) task.data.batcher = Batcher(test_batch_size=batch_size) task.data.sort_key = None batches = task.data.batches(Stage.TEST, data_source=data_source) mp = torch.multiprocessing.get_context("spawn") with torch.no_grad(): results = mp.Queue() logits_writer = LogitsWriter(results, output_path, use_gzip, output_columns, ndigits_precision) logits_writer_ctx = torch.multiprocessing.spawn(logits_writer.run, join=False) for (raw_batch, tensor_dict) in batches: raw_input_tuple = (dict_zip(*raw_batch, value_only=True) if dump_raw_input else ()) model_inputs = task.model.arrange_model_inputs(tensor_dict) model_outputs = task.model(*model_inputs) results.put((raw_input_tuple, model_outputs)) results.put((None, None)) results.close() results.join_thread() logits_writer_ctx.join()
def get_logits( snapshot_path: str, use_cuda_if_available: bool, output_path: Optional[str] = None, test_path: Optional[str] = None, field_names: Optional[List[str]] = None, dump_raw_input: bool = False, ): _set_cuda(use_cuda_if_available) task, train_config, _traing_state = load(snapshot_path) print(f"Successfully loaded model from {snapshot_path}") print(f"Model on GPU? {next(task.model.parameters()).is_cuda}") if isinstance(task, NewTask): task.model.eval() data_source = _get_data_source( test_path, train_config.task.data, field_names, task ) task.data.batcher = Batcher() task.data.sort_key = None batches = task.data.batches(Stage.TEST, data_source=data_source) with PathManager.open( output_path, "w", encoding="utf-8" ) as fout, torch.no_grad(): for (raw_batch, tensor_dict) in batches: raw_input_tuple = ( dict_zip(*raw_batch, value_only=True) if dump_raw_input else () ) model_inputs = task.model.arrange_model_inputs(tensor_dict) model_outputs = task.model(*model_inputs) if isinstance(model_outputs, tuple): model_outputs_tuple = tuple(m.tolist() for m in model_outputs) for row in zip(*raw_input_tuple, *model_outputs_tuple): dump_row = "\t".join(json.dumps(r) for r in row) fout.write(f"{dump_row}\n") elif isinstance(model_outputs, torch.Tensor): model_outputs_list = model_outputs.tolist() for row in zip(*raw_input_tuple, model_outputs_list): dump_row = "\t".join(json.dumps(r) for r in row) fout.write(f"{dump_row}\n") else: raise Exception( "Expecting tuple or torchTensor types for model_outputs" )
def get_logits( snapshot_path: str, use_cuda_if_available: bool, output_path: Optional[str] = None, test_path: Optional[str] = None, field_names: Optional[List[str]] = None, ): _set_cuda(use_cuda_if_available) task, train_config = load(snapshot_path) print(f"Successfully loaded model from {snapshot_path}") print(f"Model on GPU? {next(task.model.parameters()).is_cuda}") if isinstance(task, NewTask): task.model.eval() data_source = _get_data_source( test_path, getattr(train_config.task.data, "source", None), field_names, task, ) task.data.batcher = Batcher() task.data.sort_key = None batches = task.data.batches(Stage.TEST, data_source=data_source) with open(output_path, "w", encoding="utf-8") as fout, torch.no_grad(): for (_, tensor_dict) in batches: model_inputs = task.model.arrange_model_inputs(tensor_dict) model_outputs = task.model(*model_inputs) if isinstance(model_outputs, tuple): model_outputs_list = [m.tolist() for m in model_outputs] for row in zip(*model_outputs_list): # row is a tuple of lists dump_row = "\t".join(json.dumps(r) for r in row) fout.write(f"{dump_row}\n") elif isinstance(model_outputs, torch.Tensor): model_outputs_list = model_outputs.tolist() for row in zip(model_outputs_list): fout.write(f"{json.dumps(row)}\n") else: raise Exception( "Expecting tuple or torchTensor types for model_outputs" )
def get_logits( snapshot_path: str, use_cuda_if_available: bool, output_path: Optional[str] = None, test_path: Optional[str] = None, field_names: Optional[List[str]] = None, ): _set_cuda(use_cuda_if_available) task, train_config = load(snapshot_path) if isinstance(task, NewTask): task.model.eval() data_source = _get_data_source(test_path, field_names, task) task.data.batcher = Batcher() batches = task.data.batches(Stage.TEST, data_source=data_source) results = [] for batch in batches: model_inputs = task.model.arrange_model_inputs(batch) model_outputs = task.model(*model_inputs) MetricReporter.aggregate_data(results, model_outputs) with open(output_path, "w", encoding="utf-8") as fout: for row in results: fout.write(f"{row}\n")
def get_logits( snapshot_path: str, use_cuda_if_available: bool, output_path: Optional[str] = None, test_path: Optional[str] = None, field_names: Optional[List[str]] = None, dump_raw_input: bool = False, batch_size: int = 16, ndigits_precision: int = 0, output_columns: Optional[List[int]] = None, use_gzip: bool = False, device_id: int = 0, fp16: bool = False, ): _set_cuda(use_cuda_if_available, device_id) task, train_config, _training_state = load(snapshot_path) print(f"Successfully loaded model from {snapshot_path}") print(f"Model on GPU? {next(task.model.parameters()).is_cuda}") print(f"CUDA device id: {torch.cuda.current_device()}") if isinstance(task, NewTask): task.model.eval() if fp16: task.model.half() data_source = _get_data_source(test_path, train_config.task.data, field_names, task) task.data.batcher = Batcher(test_batch_size=batch_size) task.data.sort_key = None batches = task.data.batches(Stage.TEST, data_source=data_source) mp = torch.multiprocessing.get_context("spawn") with torch.no_grad(): results = mp.SimpleQueue() logits_writer = LogitsWriter(results, output_path, use_gzip, ndigits_precision) logits_writer_ctx = torch.multiprocessing.spawn(logits_writer.run, join=False) for (raw_batch, tensor_dict) in batches: raw_input_tuple = (dict_zip(*raw_batch, value_only=True) if dump_raw_input else ()) model_inputs = task.model.arrange_model_inputs(tensor_dict) model_outputs = task.model(*model_inputs) # multi-encoder output if isinstance(model_outputs, tuple): # prevent breaking behaviour in default case output_columns = (range(len(model_outputs)) if not output_columns else output_columns) model_outputs = tuple(m.tolist() for i, m in enumerate(model_outputs) if i in output_columns) # single encoder output elif isinstance(model_outputs, list): model_outputs = model_outputs.tolist() else: raise Exception( "Expecting tuple or tensor types for model_outputs") results.put((raw_input_tuple, model_outputs)) results.put((None, None)) logits_writer_ctx.join() print( f"Finished logits generation for file {test_path} with output {output_path}" )