def export_test_data(output_dir, model, embed_const, query_pos_us):
im = Image.open(os.path.join(output_dir, "000000039769.jpg"))
# standard PyTorch mean-std input image normalization
transform = T.Compose([
T.Resize(800),
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# for output bounding box post-processing
def box_cxcywh_to_xyxy(x):
x_c, y_c, w, h = x.unbind(1)
b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w),
(y_c + 0.5 * h)]
return torch.stack(b, dim=1)
def rescale_bboxes(out_bbox, size):
img_w, img_h = size
b = box_cxcywh_to_xyxy(out_bbox)
b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
return b
img_preprocessed = transform(im).unsqueeze(0)
pred_logits, pred_boxes = model(img_preprocessed, embed_const,
query_pos_us)
serialize_tensors(
os.path.join(output_dir, "expected.bin"), {
"input_0": img_preprocessed.numpy(),
"input_embed_const": embed_const.numpy(),
"input_query_pos_us": query_pos_us.numpy(),
"output_logits": pred_logits.numpy(),
"output_boxes": pred_boxes.numpy()
})
def main():
parser = argparse.ArgumentParser()
parser.add_argument("src", help="Source onnx model")
parser.add_argument("dst", help="Destination onnx model")
parser.add_argument(
"inputshapes",
help=
"Input tensor names and shapes of subgraph. Example: 'tensor3=10x256,tensor4=-1x-1'"
)
parser.add_argument(
"outputshapes",
help=
"Output tensor names and shapes of subgraph. Example: 'tensor3=10x256,tensor4=-1x-1'"
)
parser.add_argument(
"--test", help="Outputs test case file. Also needs --test-shapes")
parser.add_argument("--test-shapes")
args = parser.parse_args()
input_shapes = parse_name_shapes(args.inputshapes)
output_shapes = parse_name_shapes(args.outputshapes)
model = onnx.load_model(args.src)
process_model(model, input_shapes, output_shapes)
try:
onnx.checker.check_model(model)
except Exception as ex:
print("Checker error")
print(ex)
onnx.save_model(model, args.dst)
print("Saved model")
if args.test:
test_shapes = parse_name_shapes(args.test_shapes)
test_io_case = make_test_io(args.dst, output_shapes, test_shapes)
serialize_tensors(args.test, test_io_case)
print("Saved test io case")
def export_test_case(output_dir, model):
example_input = torch.randn((1, 3, 224, 224))
with torch.no_grad():
example_output = model(example_input)
serialize_tensors(os.path.join(output_dir, "expected.bin"), {
"input": example_input.numpy(),
"output": example_output.numpy(),
})
def main():
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
os.makedirs(output_dir, exist_ok=True)
torch.set_grad_enabled(False)
download_sample_image(os.path.join(output_dir, "000000039769.jpg"), 'http://images.cocodataset.org/val2017/000000039769.jpg')
state_dict = torch.hub.load_state_dict_from_url(
url='https://dl.fbaipublicfiles.com/detr/detr_demo-da2a99e9.pth',
map_location='cpu', check_hash=True)
state_dict_onnx = {k:v.float() for k, v in state_dict.items()}
H, W = 25, 34
query_pos_us = state_dict_onnx["query_pos"].unsqueeze(1)
embed_const = torch.cat([
state_dict_onnx["col_embed"][:W].unsqueeze(0).repeat(H, 1, 1),
state_dict_onnx["row_embed"][:H].unsqueeze(1).repeat(1, W, 1),
], dim=-1).flatten(0, 1).unsqueeze(1)
del state_dict_onnx["row_embed"]
del state_dict_onnx["col_embed"]
del state_dict_onnx["query_pos"]
detr_onnx = DETROnnx(num_classes=91)
detr_onnx.load_state_dict(state_dict_onnx)
detr_onnx.eval()
img_preprocessed = torch.zeros(1, 3, 800, 1066)
onnx_path = f"{output_dir}/unoptimized_model.onnx"
torch.onnx.export(detr_onnx, (img_preprocessed, embed_const, query_pos_us), onnx_path,
verbose=True,
input_names=["input_0", "input_embed_const", "input_query_pos_us"],
output_names=["output_logits", "output_boxes"], opset_version=10)
subprocess.check_call(["python", "-m", "webdnn.optimize_model", onnx_path, os.path.join(output_dir)])
serialize_tensors(os.path.join(output_dir, "embedding.bin"), {
"col_embed": state_dict["col_embed"].float().numpy(),
"row_embed": state_dict["row_embed"].float().numpy(),
"query_pos_us": query_pos_us.numpy()})
export_test_data(output_dir, detr_onnx, embed_const, query_pos_us)
def dump_expected(directory, arrays_dict):
serialize_tensors(directory + "/expected.bin", arrays_dict)