|
|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn as nn |
|
|
|
|
|
from torch.utils.mobile_optimizer import optimize_for_mobile |
|
|
|
|
|
|
|
|
import models |
|
|
import models |
|
|
from models.experimental import attempt_load |
|
|
from models.experimental import attempt_load |
|
|
|
|
|
|
|
|
print(f'\n{prefix} starting export with torch {torch.__version__}...') |
|
|
print(f'\n{prefix} starting export with torch {torch.__version__}...') |
|
|
f = opt.weights.replace('.pt', '.torchscript.pt') # filename |
|
|
f = opt.weights.replace('.pt', '.torchscript.pt') # filename |
|
|
ts = torch.jit.trace(model, img, strict=False) |
|
|
ts = torch.jit.trace(model, img, strict=False) |
|
|
|
|
|
ts = optimize_for_mobile(ts) # https://pytorch.org/tutorials/recipes/script_optimized.html |
|
|
ts.save(f) |
|
|
ts.save(f) |
|
|
print(f'{prefix} export success, saved as {f}') |
|
|
print(f'{prefix} export success, saved as {f}') |
|
|
except Exception as e: |
|
|
except Exception as e: |