diff --git a/export.py b/export.py index c2d498d..2b27b4a 100644 --- a/export.py +++ b/export.py @@ -6,6 +6,7 @@ sys.path.append('./') # to run '$ python *.py' files in subdirectories import torch import torch.nn as nn +from torch.utils.mobile_optimizer import optimize_for_mobile import models from models.experimental import attempt_load, End2End @@ -75,6 +76,17 @@ if __name__ == '__main__': except Exception as e: print('TorchScript export failure: %s' % e) + # TorchScript-Lite export + try: + print('\nStarting TorchScript-Lite export with torch %s...' % torch.__version__) + f = opt.weights.replace('.pt', '.torchscript.ptl') # filename + ts = torch.jit.trace(model, img, strict=False) + ts = optimize_for_mobile(ts) + ts._save_for_lite_interpreter(f) + print('TorchScript-Lite export success, saved as %s' % f) + except Exception as e: + print('TorchScript-Lite export failure: %s' % e) + # ONNX export try: import onnx