Commit 60205f80 authored by Yuxin Wu's avatar Yuxin Wu Committed by Facebook GitHub Bot
Browse files

allow converter to dump torchscript; support torchscript GPU inference

Reviewed By: rbgirshick

Differential Revision: D21364163

fbshipit-source-id: 6d83968b483f91df976939d8682031a5c60dd271
parent 1e21fa4e
......@@ -15,7 +15,13 @@ from .caffe2_inference import ProtobufDetectionModel
from .caffe2_modeling import META_ARCH_CAFFE2_EXPORT_TYPE_MAP, convert_batched_inputs_to_c2_format
from .shared import get_pb_arg_vali, get_pb_arg_vals, save_graph
__all__ = ["add_export_config", "export_caffe2_model", "Caffe2Model", "export_onnx_model"]
__all__ = [
"add_export_config",
"export_caffe2_model",
"Caffe2Model",
"export_onnx_model",
"Caffe2Tracer",
]
def add_export_config(cfg):
......@@ -47,7 +53,8 @@ class Caffe2Tracer:
3. complicated pre/post processing
This class provides a traceable version of a detectron2 model by:
1. Rewrite parts of the model using ops in caffe2
1. Rewrite parts of the model using ops in caffe2. Note that some ops do
not have GPU implementation.
2. Define the inputs "after pre-processing" as inputs to the model
3. Remove post-processing and produce raw layer outputs
......@@ -59,8 +66,6 @@ class Caffe2Tracer:
model to different deployment formats.
The class currently only supports models using builtin meta architectures.
Experimental. Don't use.
"""
def __init__(self, cfg, model, inputs):
......@@ -127,7 +132,7 @@ class Caffe2Tracer:
logger = logging.getLogger(__name__)
logger.info("Tracing the model with torch.jit.trace ...")
with torch.no_grad():
return torch.jit.trace(model, (inputs,))
return torch.jit.trace(model, (inputs,), optimize=True)
def export_caffe2_model(cfg, model, inputs):
......
......@@ -164,9 +164,6 @@ class Caffe2RPN(Caffe2Compatible, rpn.RPN):
features = [features[f] for f in self.in_features]
objectness_logits_pred, anchor_deltas_pred = self.rpn_head(features)
# TODO is the needed?
# objectness_logits_pred = [t.sigmoid() for t in objectness_logits_pred]
assert isinstance(images, ImageList)
if self.tensor_mode:
im_info = images.image_sizes
......
......@@ -12,3 +12,10 @@ target_link_libraries(
caffe2_mask_rcnn
"${TORCH_LIBRARIES}" gflags glog ${OpenCV_LIBS})
set_property(TARGET caffe2_mask_rcnn PROPERTY CXX_STANDARD 14)
add_executable(torchscript_traced_mask_rcnn torchscript_traced_mask_rcnn.cpp)
target_link_libraries(
torchscript_traced_mask_rcnn
"${TORCH_LIBRARIES}" ${OpenCV_LIBS})
set_property(TARGET torchscript_traced_mask_rcnn PROPERTY CXX_STANDARD 14)
......@@ -2,13 +2,14 @@
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import argparse
import os
import onnx
import torch
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import build_detection_test_loader
from detectron2.evaluation import COCOEvaluator, inference_on_dataset, print_csv_format
from detectron2.export import add_export_config, export_caffe2_model
from detectron2.export import Caffe2Tracer, add_export_config
from detectron2.modeling import build_model
from detectron2.utils.logger import setup_logger
......@@ -28,10 +29,16 @@ def setup_cfg(args):
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Convert a model to Caffe2")
parser = argparse.ArgumentParser(description="Convert a model using caffe2 tracing.")
parser.add_argument(
"--format",
choices=["caffe2", "onnx", "torchscript"],
help="output format",
default="caffe2",
)
parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file")
parser.add_argument("--run-eval", action="store_true")
parser.add_argument("--output", help="output directory for the converted caffe2 model")
parser.add_argument("--output", help="output directory for the converted model")
parser.add_argument(
"opts",
help="Modify config options using the command-line",
......@@ -41,6 +48,7 @@ if __name__ == "__main__":
args = parser.parse_args()
logger = setup_logger()
logger.info("Command line arguments: " + str(args))
os.makedirs(args.output, exist_ok=True)
cfg = setup_cfg(args)
......@@ -53,13 +61,35 @@ if __name__ == "__main__":
first_batch = next(iter(data_loader))
# convert and save caffe2 model
caffe2_model = export_caffe2_model(cfg, torch_model, first_batch)
caffe2_model.save_protobuf(args.output)
# draw the caffe2 graph
caffe2_model.save_graph(os.path.join(args.output, "model.svg"), inputs=first_batch)
tracer = Caffe2Tracer(cfg, torch_model, first_batch)
if args.format == "caffe2":
caffe2_model = tracer.export_caffe2()
caffe2_model.save_protobuf(args.output)
# draw the caffe2 graph
caffe2_model.save_graph(os.path.join(args.output, "model.svg"), inputs=first_batch)
elif args.format == "onnx":
onnx_model = tracer.export_onnx()
onnx.save(onnx_model, os.path.join(args.output, "model.onnx"))
elif args.format == "torchscript":
script_model = tracer.export_torchscript()
script_model.save(os.path.join(args.output, "model.ts"))
# Recursively print IR of all modules
with open(os.path.join(args.output, "model_ts_IR.txt"), "w") as f:
try:
f.write(script_model._actual_script_module._c.dump_to_str(True, False, False))
except AttributeError:
pass
# Print IR of the entire graph (all submodules inlined)
with open(os.path.join(args.output, "model_ts_IR_inlined.txt"), "w") as f:
f.write(str(script_model.inlined_graph))
# Print the model structure in pytorch style
with open(os.path.join(args.output, "model.txt"), "w") as f:
f.write(str(script_model))
# run evaluation with the converted model
if args.run_eval:
assert args.format == "caffe2", "Python inference in other format is not yet supported."
dataset = cfg.DATASETS.TEST[0]
data_loader = build_detection_test_loader(cfg, dataset)
# NOTE: hard-coded evaluator. change to the evaluator for your dataset
......
......@@ -9,7 +9,7 @@
using namespace std;
// Experimental. Don't use.
// experimental. don't use
int main(int argc, const char* argv[]) {
if (argc != 3) {
return 1;
......@@ -19,26 +19,25 @@ int main(int argc, const char* argv[]) {
torch::autograd::AutoGradMode guard(false);
auto module = torch::jit::load(argv[1]);
assert(module.buffers().size() > 0);
// Assume that the entire model is on the same device.
// We just put input to this device.
auto device = (*begin(module.buffers())).device();
cv::Mat input_img = cv::imread(image_file, cv::IMREAD_COLOR);
const int height = input_img.rows;
const int width = input_img.cols;
// FPN models require divisibility of 32
assert(height % 32 == 0 && width % 32 == 0);
const int batch = 1;
const int channels = 3;
auto input = torch::empty({1, channels, height, width});
float* ptr = input.data_ptr<float>();
// HWC to CHW
for (int c = 0; c < 3; ++c) {
for (int i = 0; i < height * width; ++i) {
ptr[c * height * width + i] =
static_cast<float>(input_img.data[3 * i + c]);
}
}
auto input = torch::from_blob(
input_img.data, {1, height, width, channels}, torch::kUInt8);
// NHWC to NCHW
input = input.to(device, torch::kFloat).permute({0, 3, 1, 2}).contiguous();
float im_info_data[] = {height * 1.0f, width * 1.0f, 1.0f};
auto im_info = torch::from_blob(im_info_data, {1, 3});
std::array<float, 3> im_info_data{height * 1.0f, width * 1.0f, 1.0f};
auto im_info = torch::from_blob(im_info_data.data(), {1, 3}).to(device);
// run the network
auto output = module.forward({std::make_tuple(input, im_info)});
......
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