Make a detectron2 model traceable with caffe2 style. An original detectron2 model may not be traceable, or cannot be deployed directly after being traced, due to some reasons: 1. control flow in some ops 2. custom ops 3. complicated pre/post processing This class prov
| 41 | |
| 42 | |
| 43 | class Caffe2Tracer: |
| 44 | """ |
| 45 | Make a detectron2 model traceable with caffe2 style. |
| 46 | |
| 47 | An original detectron2 model may not be traceable, or |
| 48 | cannot be deployed directly after being traced, due to some reasons: |
| 49 | |
| 50 | 1. control flow in some ops |
| 51 | 2. custom ops |
| 52 | 3. complicated pre/post processing |
| 53 | |
| 54 | This class provides a traceable version of a detectron2 model by: |
| 55 | |
| 56 | 1. Rewrite parts of the model using ops in caffe2. Note that some ops do |
| 57 | not have GPU implementation. |
| 58 | 2. Define the inputs "after pre-processing" as inputs to the model |
| 59 | 3. Remove post-processing and produce raw layer outputs |
| 60 | |
| 61 | More specifically about inputs: all builtin models take two input tensors. |
| 62 | |
| 63 | 1. NCHW float "data" which is an image (usually in [0, 255]) |
| 64 | 2. Nx3 float "im_info", each row of which is (height, width, 1.0) |
| 65 | |
| 66 | After making a traceable model, the class provide methods to export such a |
| 67 | model to different deployment formats. |
| 68 | |
| 69 | The class currently only supports models using builtin meta architectures. |
| 70 | """ |
| 71 | |
| 72 | def __init__(self, cfg, model, inputs): |
| 73 | """ |
| 74 | Args: |
| 75 | cfg (CfgNode): a detectron2 config, with extra export-related options |
| 76 | added by :func:`add_export_config`. |
| 77 | model (nn.Module): a model built by |
| 78 | :func:`detectron2.modeling.build_model`. Weights have to be already |
| 79 | loaded to this model. |
| 80 | inputs: sample inputs that the given model takes for inference. |
| 81 | Will be used to trace the model. Random input with no detected objects |
| 82 | will not work if the model has data-dependent control flow (e.g., R-CNN). |
| 83 | """ |
| 84 | assert isinstance(cfg, CN), cfg |
| 85 | assert isinstance(model, torch.nn.Module), type(model) |
| 86 | if "EXPORT_CAFFE2" not in cfg: |
| 87 | cfg = add_export_config(cfg) # will just the defaults |
| 88 | |
| 89 | self.cfg = cfg |
| 90 | self.model = model |
| 91 | self.inputs = inputs |
| 92 | |
| 93 | def _get_traceable(self): |
| 94 | # TODO how to make it extensible to support custom models |
| 95 | C2MetaArch = META_ARCH_CAFFE2_EXPORT_TYPE_MAP[self.cfg.MODEL.META_ARCHITECTURE] |
| 96 | traceable_model = C2MetaArch(self.cfg, copy.deepcopy(self.model)) |
| 97 | traceable_inputs = traceable_model.get_caffe2_inputs(self.inputs) |
| 98 | return traceable_model, traceable_inputs |
| 99 | |
| 100 | def export_caffe2(self): |
no outgoing calls
no test coverage detected