def export(args: ExportArgs, model: Module, val_loader: DataLoader,
save_dir: str) -> None:
"""
Utility method to export the model and data
:param args : An ExportArgs object containing config for export task.
:param model: loaded model architecture to export
:param val_loader: A DataLoader for validation data
:param save_dir: Directory to store checkpoints at during exporting process
"""
exporter = ModuleExporter(model, save_dir)
# export PyTorch state dict
LOGGER.info(f"exporting pytorch in {save_dir}")
exporter.export_pytorch(use_zipfile_serialization_if_available=(
args.use_zipfile_serialization_if_available))
onnx_exported = False
for batch, data in tqdm(
enumerate(val_loader),
desc="Exporting samples",
total=args.num_samples if args.num_samples > 1 else 1,
):
if not onnx_exported:
# export onnx file using first sample for graph freezing
LOGGER.info(f"exporting onnx in {save_dir}")
exporter.export_onnx(data[0],
opset=args.onnx_opset,
convert_qat=True)
onnx_exported = True
if args.num_samples > 0:
exporter.export_samples(sample_batches=[data[0]],
sample_labels=[data[1]],
exp_counter=batch)
def test_export_batches(batch_size):
sample_batch = torch.randn(batch_size, 8)
exporter = ModuleExporter(MLPNet(), tempfile.gettempdir())
exporter.export_samples([sample_batch])