Support for mobile

[gillpeacegood]

Hi Not sure if this is bug or feature or neither, but I have the following problem with trying to save a model for lite interpreter per pytorch mobile pipeline

from torch.utils.mobile_optimizer import optimize_for_mobile
yolo = YOLOv5.load_from_yolov5('models/mymodel.pt')
yolo_script = torch.jit.script(yolo)
optimized_yolo = optimize_for_mobile(yolo_script)
optimized_yolo._save_for_lite_interpreter('models/mymodel.ptl')

And this happens

----> 8 optimized_yolo._save_for_lite_interpreter('models/mymodel.ptl')

~/opt/anaconda3/envs/axsy/lib/python3.9/site-packages/torch/jit/_script.py in _save_for_lite_interpreter(self, *args, **kwargs)
    726 
    727             """
--> 728             return self._c._save_for_mobile(*args, **kwargs)
    729 
    730         def _save_to_buffer_for_lite_interpreter(self, *args, **kwargs):

RuntimeError: __torch__ types other than custom c++ classes (__torch__.torch.classes)are not supported in lite interpreter. Workaround: instead of using arbitrary class type (class Foo()), define a pytorch class (class Foo(torch.nn.Module)). The problematic type is: __torch__.yolort.models.transform.NestedTensor

[zhiqwang]

Hi @gillpeacegood , there is a working example for mobile deployment of yolort:

## Export the TorchScript-ed model

import torch
import torchvision
from torch.utils.mobile_optimizer import optimize_for_mobile
from yolort.models import YOLO
from yolort.v5 import attempt_download

# Prepare some parameters for the exported torchscript-ed module
score_thresh = 0.25
nms_thresh = 0.45
device = torch.device("cpu")

# Downloaded from 'https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt'
model_path = "yolov5s.pt"
checkpoint_path = attempt_download(model_path)

model = YOLO.load_from_yolov5(checkpoint_path, score_thresh=score_thresh, nms_thresh=nms_thresh)

model = model.eval()
model = model.to(device)

export_scripted_path = "yolov5s_scripted.pt"
export_optimized_path = "yolov5s_scriptmodule.ptl"

scripted_model = torch.jit.script(model)
scripted_model.save(export_scripted_path)
optimized_model = optimize_for_mobile(scripted_model)
optimized_model._save_for_lite_interpreter(export_optimized_path)

See facebookresearch/playtorch#10 (comment) for more details.

[gillpeacegood]

Sorry to be stupid but I can't see how it removes the pre processing

[zhiqwang]

We provide the class for removing pre-processing directly, note that this is YOLO on the above example, not YOLOv5. See #265 (comment) and https://zhiqwang.com/yolov5-rt-stack/notebooks/how-to-align-with-ultralytics-yolov5.html for more details.

[gillpeacegood]

As okay thanks. Shame about the preprocessing though :-)

[gillpeacegood]

Anyway the error suggests that the issue could be resolved by moving NestedTestor to a pytorch class, if it works would you accept a PR on that?

[zhiqwang]

Anyway the error suggests that the issue could be resolved by moving NestedTestor to a pytorch class, if it works would you accept a PR on that?

Hi @gillpeacegood , Yep, it will be great!