Skip to content

GitLab

Commit f018d4a7 authored by Yanghan Wang's avatar Yanghan Wang Committed by Facebook GitHub Bot
Browse files

make exporting rcnn model using torchvision ops the default option 

Reviewed By: zhanghang1989

Differential Revision: D31134064

fbshipit-source-id: 825ca14477243a53f84b8521f4430a2b080324bd
parent 3ce16d73
Show whitespace changes
Inline Side-by-side
Showing with 220 additions and 156 deletions
d2go/export/torchscript.py
View file @ f018d4a7
......@@ -214,10 +214,10 @@ def tracing_adapter_wrap_load(old_f):
return new_f
@ModelExportMethodRegistry.register("torchscript")
@ModelExportMethodRegistry.register("torchscript_int8")
@ModelExportMethodRegistry.register("torchscript_mobile")
@ModelExportMethodRegistry.register("torchscript_mobile_int8")
@ModelExportMethodRegistry.register("torchscript@legacy")
@ModelExportMethodRegistry.register("torchscript_int8@legacy")
@ModelExportMethodRegistry.register("torchscript_mobile@legacy")
@ModelExportMethodRegistry.register("torchscript_mobile_int8@legacy")
class DefaultTorchscriptExport(ModelExportMethod):
@classmethod
def export(
......@@ -250,10 +250,10 @@ class DefaultTorchscriptExport(ModelExportMethod):
return load_torchscript(model_path)
@ModelExportMethodRegistry.register("torchscript@tracing")
@ModelExportMethodRegistry.register("torchscript_int8@tracing")
@ModelExportMethodRegistry.register("torchscript_mobile@tracing")
@ModelExportMethodRegistry.register("torchscript_mobile_int8@tracing")
@ModelExportMethodRegistry.register("torchscript")
@ModelExportMethodRegistry.register("torchscript_int8")
@ModelExportMethodRegistry.register("torchscript_mobile")
@ModelExportMethodRegistry.register("torchscript_mobile_int8")
class D2TorchscriptTracingExport(DefaultTorchscriptExport):
@classmethod
@tracing_adapter_wrap_export
......
d2go/modeling/meta_arch/rcnn.py
View file @ f018d4a7
......@@ -5,14 +5,16 @@
import logging
import torch
import torch.nn as nn
from caffe2.proto import caffe2_pb2
from d2go.export.api import PredictorExportConfig
from d2go.utils.export_utils import (
D2Caffe2MetaArchPreprocessFunc,
D2Caffe2MetaArchPostprocessFunc,
D2RCNNTracingWrapper,
from detectron2.export.caffe2_modeling import (
META_ARCH_CAFFE2_EXPORT_TYPE_MAP,
convert_batched_inputs_to_c2_format,
)
from detectron2.export.caffe2_modeling import META_ARCH_CAFFE2_EXPORT_TYPE_MAP
from detectron2.export.shared import get_pb_arg_vali, get_pb_arg_vals
from detectron2.modeling import GeneralizedRCNN
from detectron2.modeling.postprocessing import detector_postprocess
from detectron2.projects.point_rend import PointRendMaskHead
from detectron2.utils.registry import Registry
from mobile_cv.arch.utils import fuse_utils
......@@ -60,22 +62,11 @@ class GeneralizedRCNNPatch:
@RCNN_PREPARE_FOR_EXPORT_REGISTRY.register()
def default_rcnn_prepare_for_export(self, cfg, inputs, predictor_type):
if "torchscript" in predictor_type and "@tracing" in predictor_type:
preprocess_info = FuncInfo.gen_func_info(
D2RCNNTracingWrapper.Preprocess, params={}
)
preprocess_func = preprocess_info.instantiate()
return PredictorExportConfig(
model=D2RCNNTracingWrapper(self),
data_generator=lambda x: (preprocess_func(x),),
preprocess_info=preprocess_info,
postprocess_info=FuncInfo.gen_func_info(
D2RCNNTracingWrapper.Postprocess, params={}
),
)
if cfg.MODEL.META_ARCHITECTURE in META_ARCH_CAFFE2_EXPORT_TYPE_MAP:
if (
"@legacy" in predictor_type
or "caffe2" in predictor_type
or "onnx" in predictor_type
):
C2MetaArch = META_ARCH_CAFFE2_EXPORT_TYPE_MAP[cfg.MODEL.META_ARCHITECTURE]
c2_compatible_model = C2MetaArch(cfg, self)
......@@ -99,7 +90,19 @@ def default_rcnn_prepare_for_export(self, cfg, inputs, predictor_type):
postprocess_info=postprocess_info,
)
raise NotImplementedError("Can't determine prepare_for_tracing!")
else:
preprocess_info = FuncInfo.gen_func_info(
D2RCNNInferenceWrapper.Preprocess, params={}
)
preprocess_func = preprocess_info.instantiate()
return PredictorExportConfig(
model=D2RCNNInferenceWrapper(self),
data_generator=lambda x: (preprocess_func(x),),
preprocess_info=preprocess_info,
postprocess_info=FuncInfo.gen_func_info(
D2RCNNInferenceWrapper.Postprocess, params={}
),
)
def _apply_eager_mode_quant(cfg, model):
......@@ -278,3 +281,114 @@ def default_rcnn_prepare_for_quant_convert(self, cfg):
self.roi_heads.box_predictor.bbox_pred
)
return self
class D2Caffe2MetaArchPreprocessFunc(object):
def __init__(self, size_divisibility, device):
self.size_divisibility = size_divisibility
self.device = device
def __call__(self, inputs):
data, im_info = convert_batched_inputs_to_c2_format(
inputs, self.size_divisibility, self.device
)
return (data, im_info)
@staticmethod
def get_params(cfg, model):
fake_predict_net = caffe2_pb2.NetDef()
model.encode_additional_info(fake_predict_net, None)
size_divisibility = get_pb_arg_vali(fake_predict_net, "size_divisibility", 0)
device = get_pb_arg_vals(fake_predict_net, "device", b"cpu").decode("ascii")
return {
"size_divisibility": size_divisibility,
"device": device,
}
class D2Caffe2MetaArchPostprocessFunc(object):
def __init__(self, external_input, external_output, encoded_info):
self.external_input = external_input
self.external_output = external_output
self.encoded_info = encoded_info
def __call__(self, inputs, tensor_inputs, tensor_outputs):
encoded_info = self.encoded_info.encode("ascii")
fake_predict_net = caffe2_pb2.NetDef().FromString(encoded_info)
meta_architecture = get_pb_arg_vals(fake_predict_net, "meta_architecture", None)
meta_architecture = meta_architecture.decode("ascii")
model_class = META_ARCH_CAFFE2_EXPORT_TYPE_MAP[meta_architecture]
convert_outputs = model_class.get_outputs_converter(fake_predict_net, None)
c2_inputs = tensor_inputs
c2_results = dict(zip(self.external_output, tensor_outputs))
return convert_outputs(inputs, c2_inputs, c2_results)
@staticmethod
def get_params(cfg, model):
# NOTE: the post processing has different values for different meta
# architectures, here simply relying Caffe2 meta architecture to encode info
# into a NetDef and storing it as whole.
fake_predict_net = caffe2_pb2.NetDef()
model.encode_additional_info(fake_predict_net, None)
encoded_info = fake_predict_net.SerializeToString().decode("ascii")
# HACK: Caffe2MetaArch's post processing requires the blob name of model output,
# this information is missed for torchscript. There's no easy way to know this
# unless using NamedTuple for tracing.
external_input = ["data", "im_info"]
if cfg.MODEL.META_ARCHITECTURE == "GeneralizedRCNN":
external_output = ["bbox_nms", "score_nms", "class_nms"]
if cfg.MODEL.MASK_ON:
external_output.extend(["mask_fcn_probs"])
if cfg.MODEL.KEYPOINT_ON:
if cfg.EXPORT_CAFFE2.USE_HEATMAP_MAX_KEYPOINT:
external_output.extend(["keypoints_out"])
else:
external_output.extend(["kps_score"])
else:
raise NotImplementedError("")
return {
"external_input": external_input,
"external_output": external_output,
"encoded_info": encoded_info,
}
class D2RCNNInferenceWrapper(nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, image):
"""
This function describes what happends during the tracing. Note that the output
contains non-tensor, therefore the D2TorchscriptTracingExport must be used in
order to convert the output back from flattened tensors.
"""
inputs = [{"image": image}]
return self.model.inference(inputs, do_postprocess=False)[0]
@staticmethod
class Preprocess(object):
"""
This function describes how to covert orginal input (from the data loader)
to the inputs used during the tracing (i.e. the inputs of forward function).
"""
def __call__(self, batch):
assert len(batch) == 1, "only support single batch"
return batch[0]["image"]
class Postprocess(object):
def __call__(self, batch, inputs, outputs):
"""
This function describes how to run the predictor using exported model. Note
that `tracing_adapter_wrapper` runs the traced model under the hood and
behaves exactly the same as the forward function.
"""
assert len(batch) == 1, "only support single batch"
width, height = batch[0]["width"], batch[0]["height"]
r = detector_postprocess(outputs, height, width)
return [{"instances": r}]
d2go/utils/export_utils.py deleted 100644 → 0
View file @ 3ce16d73
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import torch.nn as nn
from caffe2.proto import caffe2_pb2
from detectron2.export.caffe2_modeling import (
META_ARCH_CAFFE2_EXPORT_TYPE_MAP,
convert_batched_inputs_to_c2_format,
)
from detectron2.export.shared import get_pb_arg_vali, get_pb_arg_vals
from detectron2.modeling.postprocessing import detector_postprocess
class D2Caffe2MetaArchPreprocessFunc(object):
def __init__(self, size_divisibility, device):
self.size_divisibility = size_divisibility
self.device = device
def __call__(self, inputs):
data, im_info = convert_batched_inputs_to_c2_format(
inputs, self.size_divisibility, self.device
)
return (data, im_info)
@staticmethod
def get_params(cfg, model):
fake_predict_net = caffe2_pb2.NetDef()
model.encode_additional_info(fake_predict_net, None)
size_divisibility = get_pb_arg_vali(fake_predict_net, "size_divisibility", 0)
device = get_pb_arg_vals(fake_predict_net, "device", b"cpu").decode("ascii")
return {
"size_divisibility": size_divisibility,
"device": device,
}
class D2Caffe2MetaArchPostprocessFunc(object):
def __init__(self, external_input, external_output, encoded_info):
self.external_input = external_input
self.external_output = external_output
self.encoded_info = encoded_info
def __call__(self, inputs, tensor_inputs, tensor_outputs):
encoded_info = self.encoded_info.encode("ascii")
fake_predict_net = caffe2_pb2.NetDef().FromString(encoded_info)
meta_architecture = get_pb_arg_vals(fake_predict_net, "meta_architecture", None)
meta_architecture = meta_architecture.decode("ascii")
model_class = META_ARCH_CAFFE2_EXPORT_TYPE_MAP[meta_architecture]
convert_outputs = model_class.get_outputs_converter(fake_predict_net, None)
c2_inputs = tensor_inputs
c2_results = dict(zip(self.external_output, tensor_outputs))
return convert_outputs(inputs, c2_inputs, c2_results)
@staticmethod
def get_params(cfg, model):
# NOTE: the post processing has different values for different meta
# architectures, here simply relying Caffe2 meta architecture to encode info
# into a NetDef and storing it as whole.
fake_predict_net = caffe2_pb2.NetDef()
model.encode_additional_info(fake_predict_net, None)
encoded_info = fake_predict_net.SerializeToString().decode("ascii")
# HACK: Caffe2MetaArch's post processing requires the blob name of model output,
# this information is missed for torchscript. There's no easy way to know this
# unless using NamedTuple for tracing.
external_input = ["data", "im_info"]
if cfg.MODEL.META_ARCHITECTURE == "GeneralizedRCNN":
external_output = ["bbox_nms", "score_nms", "class_nms"]
if cfg.MODEL.MASK_ON:
external_output.extend(["mask_fcn_probs"])
if cfg.MODEL.KEYPOINT_ON:
if cfg.EXPORT_CAFFE2.USE_HEATMAP_MAX_KEYPOINT:
external_output.extend(["keypoints_out"])
else:
external_output.extend(["kps_score"])
else:
raise NotImplementedError("")
return {
"external_input": external_input,
"external_output": external_output,
"encoded_info": encoded_info,
}
class D2RCNNTracingWrapper(nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, image):
"""
This function describes what happends during the tracing. Note that the output
contains non-tensor, therefore the D2TorchscriptTracingExport must be used in
order to convert the output back from flattened tensors.
"""
inputs = [{"image": image}]
return self.model.inference(inputs, do_postprocess=False)[0]
@staticmethod
class Preprocess(object):
"""
This function describes how to covert orginal input (from the data loader)
to the inputs used during the tracing (i.e. the inputs of forward function).
"""
def __call__(self, batch):
assert len(batch) == 1, "only support single batch"
return batch[0]["image"]
class Postprocess(object):
def __call__(self, batch, inputs, outputs):
"""
This function describes how to run the predictor using exported model. Note
that `tracing_adapter_wrapper` runs the traced model under the hood and
behaves exactly the same as the forward function.
"""
assert len(batch) == 1, "only support single batch"
width, height = batch[0]["width"], batch[0]["height"]
r = detector_postprocess(outputs, height, width)
return [{"instances": r}]
tests/modeling/test_meta_arch_rcnn.py
View file @ f018d4a7
......@@ -29,9 +29,10 @@ class TestFBNetV3MaskRCNNFP32(RCNNBaseTestCases.TemplateTestCase):
@RCNNBaseTestCases.expand_parameterized_test_export(
[
["torchscript@tracing", True],
["torchscript@legacy", True],
["torchscript", True],
["torchscript_int8@legacy", False],
["torchscript_int8", False],
["torchscript_int8@tracing", False],
]
)
def test_export(self, predictor_type, compare_match):
......@@ -59,6 +60,7 @@ class TestFBNetV3MaskRCNNQATEager(RCNNBaseTestCases.TemplateTestCase):
@RCNNBaseTestCases.expand_parameterized_test_export(
[
["torchscript_int8@legacy", False], # TODO: fix mismatch
["torchscript_int8", False], # TODO: fix mismatch
]
)
......@@ -85,6 +87,7 @@ class TestFBNetV3KeypointRCNNFP32(RCNNBaseTestCases.TemplateTestCase):
@RCNNBaseTestCases.expand_parameterized_test_export(
[
["torchscript_int8@legacy", False], # TODO: fix mismatch
["torchscript_int8", False], # TODO: fix mismatch
]
)
......@@ -133,7 +136,7 @@ class TestTorchVisionExport(unittest.TestCase):
predictor_path = convert_and_export_predictor(
cfg,
copy.deepcopy(pytorch_model),
"torchscript@tracing",
"torchscript",
tmp_dir,
data_loader,
)
......
tests/modeling/test_rcnn_export_example.py 0 → 100644
View file @ f018d4a7
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import os
import unittest
import torch
from d2go.runner.default_runner import GeneralizedRCNNRunner
from d2go.tools.exporter import main
from d2go.utils.testing.rcnn_helper import get_quick_test_config_opts
from mobile_cv.common.misc.file_utils import make_temp_directory
def maskrcnn_export_legacy_vs_new_format_example():
with make_temp_directory("export_demo") as tmp_dir:
# START_WIKI_EXAMPLE_TAG
runner = GeneralizedRCNNRunner()
cfg = runner.get_default_cfg()
cfg.merge_from_file("detectron2go://mask_rcnn_fbnetv3a_dsmask_C4.yaml")
cfg.merge_from_list(get_quick_test_config_opts())
# equivalent to running:
# exporter.par --runner GeneralizedRCNNRunner --config-file config.yaml --predictor-types torchscript tourchscript@legacy --output-dir tmp_dir
_ = main(
cfg, tmp_dir, runner, predictor_types=["torchscript@legacy", "torchscript"]
)
# the path can be fetched from the return of main, here just use hard-coded values
new_path = os.path.join(tmp_dir, "torchscript", "model.jit")
lagacy_path = os.path.join(tmp_dir, "torchscript@legacy", "model.jit")
new_model = torch.jit.load(new_path)
legacy_model = torch.jit.load(lagacy_path)
# Running inference using new format
image = torch.zeros(1, 64, 96) # chw 3D tensor
new_outputs = new_model(image) # suppose N instances are detected
# NOTE: the output are flattened tensors of the real output (which is a dict), they're
# ordered by the key in dict, which is deterministic for the given model, but it might
# be difficult to figure out just from model.jit file. The predictor_info.json from
# the same directory contains the `outputs_schema`, which indicate how the final output
# is constructed from flattened tensors.
pred_boxes = new_outputs[0] # torch.Size([N, 4])
pred_classes = new_outputs[1] # torch.Size([N])
pred_masks = new_outputs[2] # torch.Size([N, 1, Hmask, Wmask])
scores = new_outputs[3] # torch.Size([N])
# Running inference using legacy caffe2 format
data = torch.zeros(1, 1, 64, 96)
im_info = torch.tensor([[64, 96, 1.0]])
legacy_outputs = legacy_model([data, im_info])
# NOTE: the output order is determined in the order of creating the tensor during
# forward function, it's also follow the order of original Caffe2 model.
roi_bbox_nms = legacy_outputs[0] # torch.Size([N, 4])
roi_score_nms = legacy_outputs[1] # torch.Size([N])
roi_class_nms = legacy_outputs[2] # torch.Size([N])
mask_fcn_probs = legacy_outputs[3] # torch.Size([N, Cmask, Hmask, Wmask])
# relations between legacy outputs and new outputs
torch.testing.assert_allclose(pred_boxes, roi_bbox_nms)
torch.testing.assert_allclose(pred_classes, roi_class_nms)
torch.testing.assert_allclose(
pred_masks, mask_fcn_probs[:, roi_class_nms.to(torch.int64), :, :]
)
torch.testing.assert_allclose(scores, roi_score_nms)
# END_WIKI_EXAMPLE_TAG
class TestOptimizer(unittest.TestCase):
def test_maskrcnn_export_legacy_vs_new_format_example(self):
maskrcnn_export_legacy_vs_new_format_example()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment