make_assets.py 634 B

123456789101112131415161718192021
  1. import torch
  2. from torch.utils.mobile_optimizer import optimize_for_mobile
  3. from torchvision.models.detection import (
  4. fasterrcnn_mobilenet_v3_large_320_fpn,
  5. FasterRCNN_MobileNet_V3_Large_320_FPN_Weights,
  6. )
  7. print(torch.__version__)
  8. model = fasterrcnn_mobilenet_v3_large_320_fpn(
  9. weights=FasterRCNN_MobileNet_V3_Large_320_FPN_Weights.DEFAULT,
  10. box_score_thresh=0.7,
  11. rpn_post_nms_top_n_test=100,
  12. rpn_score_thresh=0.4,
  13. rpn_pre_nms_top_n_test=150,
  14. )
  15. model.eval()
  16. script_model = torch.jit.script(model)
  17. opt_script_model = optimize_for_mobile(script_model)
  18. opt_script_model.save("VisionTestApp/frcnn_mnetv3.pt")