glenn-jocher commited on
Commit
28db237
1 Parent(s): f662c18

Default optimize_for_mobile() on TorchScript models (#2908)

Browse files

Per https://pytorch.org/tutorials/recipes/script_optimized.html this should improve performance on torchscript models (and maybe coreml models also since coremltools operates on a torchscript model input, though this still requires testing).

Files changed (1) hide show
  1. models/export.py +2 -0
models/export.py CHANGED
@@ -12,6 +12,7 @@ sys.path.append('./') # to run '$ python *.py' files in subdirectories
12
 
13
  import torch
14
  import torch.nn as nn
 
15
 
16
  import models
17
  from models.experimental import attempt_load
@@ -65,6 +66,7 @@ if __name__ == '__main__':
65
  print(f'\n{prefix} starting export with torch {torch.__version__}...')
66
  f = opt.weights.replace('.pt', '.torchscript.pt') # filename
67
  ts = torch.jit.trace(model, img, strict=False)
 
68
  ts.save(f)
69
  print(f'{prefix} export success, saved as {f}')
70
  except Exception as e:
 
12
 
13
  import torch
14
  import torch.nn as nn
15
+ from torch.utils.mobile_optimizer import optimize_for_mobile
16
 
17
  import models
18
  from models.experimental import attempt_load
 
66
  print(f'\n{prefix} starting export with torch {torch.__version__}...')
67
  f = opt.weights.replace('.pt', '.torchscript.pt') # filename
68
  ts = torch.jit.trace(model, img, strict=False)
69
+ ts = optimize_for_mobile(ts) # https://pytorch.org/tutorials/recipes/script_optimized.html
70
  ts.save(f)
71
  print(f'{prefix} export success, saved as {f}')
72
  except Exception as e: