Skip to content

Exporter

TorchScriptExporter

class fastnn.exporting.TorchScriptExporter(model, dataloader, use_gpu=False)

Exporter Class for turning torch models into torch script models. Requires a torch Module and corresponding DataLoader populated with complete data samples. This means the DataLoader has atleast one full batch sample that can be used as input to the torch Module in order to run torch.jit.trace()

Usage:

>>> exporter = TorchScriptExporter(model=BertQAModule(...), dataloader=dataloader, use_gpu=True)
>>> exporter.export()

Parameters:

  • model - Torch model/module, preferably from fastnn.nn
  • dataloader - Torch DataLoader that corresponds to the model
  • use_gpu - Bool for using gpu or cpu. If set True but no gpu devices available, model will default to using cpu
export(self)

Traces pytorch model and returns ScriptModule model

serialize(self, file_path)

Serialize and save model

  • file_path - String file path to save serialized torchscript model