def batch_predict_caffe2_model(
pytext_model_file: str,
caffe2_model_file: str,
db_type: str = CAFFE2_DB_TYPE,
data_source: Optional[DataSource] = None,
use_cuda=False,
task: Optional[NewTask] = None,
train_config: Optional[PyTextConfig] = None,
cache_size: int = 0,
):
"""
Gets predictions from caffe2 model from a batch of examples.
Args:
pytext_model_file: Path to pytext model file (required if task and
training config is not specified)
caffe2_model_file: Path to caffe2 model file
db_type: DB type to use for caffe2
data_source: Data source for test examples
use_cuda: Whether to turn on cuda processing
task: The pytext task object
train_config: The pytext training config
cache_size: The LRU cache size to use for prediction. 0 = no cache,
-1 = boundless cache, [1, inf) = size of cache
"""
logging.info(f"Loading data processing config from {pytext_model_file}")
_set_cuda(use_cuda)
if task is None or train_config is None:
task, train_config, _ = load(pytext_model_file)
data_source = data_source or task.data.data_source
logging.info("Loading Caffe2 model")
predictor = create_predictor(train_config, caffe2_model_file, db_type,
task, cache_size)
logging.info(f"Model loaded, start testing")
predictions = [predictor(example) for example in data_source.test]
return predictions
def test_model_from_snapshot_path(
snapshot_path: str,
use_cuda_if_available: bool,
test_path: Optional[str] = None,
metric_channels: Optional[List[Channel]] = None,
test_out_path: str = "",
field_names: Optional[List[str]] = None,
):
_set_cuda(use_cuda_if_available)
task, train_config = load(snapshot_path)
for mc in metric_channels or []:
task.metric_reporter.add_channel(mc)
# Overwrite the test output path because you might not have permission to
# write to the original test output path that was created when model was trained.
if test_out_path:
if hasattr(task.metric_reporter, "output_path"):
task.metric_reporter.output_path = test_out_path
for channel in task.metric_reporter.channels:
if hasattr(channel, "file_path"):
channel.file_path = test_out_path
else:
test_out_path = train_config.task.metric_reporter.output_path
if isinstance(task, (NewTask, NewDisjointMultitask)):
data_source = _get_data_source(
test_path,
getattr(train_config.task.data, "source", None),
field_names,
task,
)
test_results = task.test(data_source)
else:
if not test_path:
test_path = train_config.task.data_handler.test_path
test_results = task.test(test_path)
return test_results, test_out_path, metric_channels
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, _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 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 prepare_task(
config: PyTextConfig,
dist_init_url: str = None,
device_id: int = 0,
rank: int = 0,
world_size: int = 1,
summary_writer: Optional[SummaryWriter] = None,
) -> Task:
if dist_init_url and world_size > 1:
dist_init(rank, world_size, dist_init_url)
print("\nParameters: {}\n".format(config))
_set_cuda(config.use_cuda_if_available, device_id, world_size)
if config.load_snapshot_path and os.path.isfile(config.load_snapshot_path):
task = load(config.load_snapshot_path)
else:
task = create_task(config.task)
if summary_writer:
task.metric_reporter.add_channel(
TensorBoardChannel(summary_writer=summary_writer))
return task
def batch_predict(model_file: str, examples: List[Dict[str, Any]]):
task, train_config = load(model_file)
return task.predict(examples)
def export_saved_model_to_torchscript(saved_model_path: str, path: str) -> None:
task, train_config = load(saved_model_path)
task.torchscript_export(task.model, path)
def export_saved_model_to_torchscript(saved_model_path: str,
path: str,
quantize: bool = False) -> None:
task, train_config = load(saved_model_path)
task.torchscript_export(task.model, path, quantize)
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}"
)
def prepare_task(
config: PyTextConfig,
dist_init_url: str = None,
device_id: int = 0,
rank: int = 0,
world_size: int = 1,
metric_channels: Optional[List[Channel]] = None,
metadata: CommonMetadata = None,
) -> Tuple[Task_Deprecated, TrainingState]:
if world_size > 1 and config.random_seed is None:
msg = (
"Must set random seed when using world_size > 1, so that parameters have "
"same initialization across workers.")
raise ValueError(msg)
if rank == 0:
print("\nParameters: {}\n".format(config), flush=True)
_set_cuda(config.use_cuda_if_available, device_id, world_size)
_set_fp16(config.use_fp16, rank)
_set_distributed(
rank,
world_size,
dist_init_url,
device_id,
config.gpu_streams_for_distributed_training,
)
if config.random_seed is not None:
set_random_seeds(config.random_seed, config.use_deterministic_cudnn)
training_state = None
if config.auto_resume_from_snapshot:
# if there are existing checkpoints, resume from the latest one
latest_snapshot_path = get_latest_checkpoint_path(
os.path.dirname(config.save_snapshot_path))
if latest_snapshot_path:
config.load_snapshot_path = latest_snapshot_path
if config.load_snapshot_path:
assert PathManager.isfile(config.load_snapshot_path)
if config.use_config_from_snapshot:
task, _, training_state = load(config.load_snapshot_path,
rank=rank,
world_size=world_size)
else:
task, _, training_state = load(
config.load_snapshot_path,
overwrite_config=config,
rank=rank,
world_size=world_size,
)
if training_state:
training_state.rank = rank
else:
task = create_task(config.task,
metadata=metadata,
rank=rank,
world_size=world_size)
for mc in metric_channels or []:
task.metric_reporter.add_channel(mc)
return task, training_state
def export_saved_model_to_torchscript(
saved_model_path: str, path: str, **kwargs
) -> None:
task, train_config, _training_state = load(saved_model_path)
task.torchscript_export(task.model, path, **kwargs)
def export_saved_model_to_torchscript(
saved_model_path: str, path: str, export_config: ExportConfig
) -> None:
task, train_config, _training_state = load(saved_model_path)
task.torchscript_export(task.model, path, False, 1, export_config=export_config)
def export_saved_model_to_torchscript(saved_model_path: str, path: str,
export_config: ExportConfig) -> None:
cuda.CUDA_ENABLED = False
precision.FP16_ENABLED = False
task, train_config, _training_state = load(saved_model_path)
task.torchscript_export(task.model, path, export_config=export_config)