Torvision model multiprocessing : pickle.PickleError during inference

Hello,

I’m trying to run inference using Pytorch Multiprocessing and torchvision fasterrcnn model. Unfortunately have some trouble using it. I would like some help on this please.

Here is the mp part:

        processes = []
        
        for data_dict in list_data_dicts:
            # defects_coordinates = detector.detection_pipeline(args.imgpath)
            p = Process(target=self.pipe, args=(data_dict, ))
            p.start()
            processes.append(p)
        
        for process in processes:
            process.join()

My model is loaded as follow:

def load_model_weights(self):
        model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
        in_features = model.roi_heads.box_predictor.cls_score.in_features
        model.roi_heads.box_predictor = (
            torchvision.models.detection.faster_rcnn.FastRCNNPredictor(in_features, self.num_classes)
        )
        checkpoint = torch.load(self.model_path, map_location=self.device)
        model.load_state_dict(checkpoint['model_state_dict'])
        return model

here is my import:

from torch.multiprocessing import Pool, Process, set_start_method
#torch.backends.cudnn.benchmark = True
try:
    set_start_method('spawn')
except RuntimeError:
    pass

And finally, the error:

Traceback (most recent call last):
  File "run_detector.py", line 57, in <module>
    defects_coordinates = detector.detection_pipeline(imgs)
  File "/data/fis/pytorch-rcnn/inference/defects_detector.py", line 70, in detection_pipeline
    p.start()
  File "/home/ansible/miniconda3/lib/python3.7/multiprocessing/process.py", line 112, in start
    self._popen = self._Popen(self)
  File "/home/ansible/miniconda3/lib/python3.7/multiprocessing/context.py", line 223, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/home/ansible/miniconda3/lib/python3.7/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/home/ansible/miniconda3/lib/python3.7/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/home/ansible/miniconda3/lib/python3.7/multiprocessing/popen_fork.py", line 20, in __init__
    self._launch(process_obj)
  File "/home/ansible/miniconda3/lib/python3.7/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/home/ansible/miniconda3/lib/python3.7/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
  File "/data/fis/pytorch-rcnn/env-rcnn/lib/python3.7/site-packages/torch/jit/__init__.py", line 1648, in __getstate__
    "Mixed serialization of script and non-script modules is not supported. " +
_pickle.PickleError: ScriptModules cannot be deepcopied using copy.deepcopy or saved using torch.save. Mixed serialization of script and non-script modules is not supported. For purely script modules use my_script_module.save(<filename>) instead.

If someone already know how to handle this, I would help me a lot.
Thank you,

Medhy

It works only if I reload my model for each process (In the self.pipe), but I want to avoid doing this.
Then I also try to copy.deepcopy the model, but this one doesn’t work too.