| | |
|---|
| | | try: |
|---|
| | | import torch_blade |
|---|
| | | except Exception as e: |
|---|
| | | print(f"failed to load torch_blade: {e}") |
|---|
| | | print(f"Warning, if you are exporting bladedisc, please install it and try it again: pip install -U torch_blade\n") |
|---|
| | | |
|---|
| | | |
|---|
| | | def export(model, data_in=None, quantize: bool = False, opset_version: int = 14, type='onnx', **kwargs): |
|---|
| | |
|---|
| | | ) |
|---|
| | | elif type == 'torchscripts': |
|---|
| | | device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|---|
| | | print("Exporting torchscripts on device {}".format(device)) |
|---|
| | | _torchscripts( |
|---|
| | | m, |
|---|
| | | path=export_dir, |
|---|
| | |
|---|
| | | model.encoder = _bladedisc_opt(model.encoder, input_data[:2]) |
|---|
| | | model.decoder = _bladedisc_opt(model.decoder, tuple(decoder_inputs)) |
|---|
| | | model_script = torch.jit.trace(model, input_data) |
|---|
| | | model_script.save(os.path.join(path, f"{model.export_name}.torchscripts")) |
|---|
| | | model_script.save(os.path.join(path, f"{model.export_name}_blade.torchscripts")) |
|---|
