Spaces:
No application file
No application file
File size: 3,496 Bytes
b26e93d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
"""
D-FINE: Redefine Regression Task of DETRs as Fine-grained Distribution Refinement
Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
---------------------------------------------------------------------------------
Modified from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright (c) 2023 lyuwenyu. All Rights Reserved.
"""
import os
import sys
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
import torch
import torch.nn as nn
from src.core import YAMLConfig
def main(
args,
):
"""main"""
cfg = YAMLConfig(args.config, resume=args.resume)
if "HGNetv2" in cfg.yaml_cfg:
cfg.yaml_cfg["HGNetv2"]["pretrained"] = False
if args.resume:
checkpoint = torch.load(args.resume, map_location="cpu")
if "ema" in checkpoint:
state = checkpoint["ema"]["module"]
else:
state = checkpoint["model"]
# NOTE load train mode state -> convert to deploy mode
cfg.model.load_state_dict(state)
else:
# raise AttributeError('Only support resume to load model.state_dict by now.')
print("not load model.state_dict, use default init state dict...")
class Model(nn.Module):
def __init__(
self,
) -> None:
super().__init__()
self.model = cfg.model.deploy()
self.postprocessor = cfg.postprocessor.deploy()
def forward(self, images, orig_target_sizes):
outputs = self.model(images)
outputs = self.postprocessor(outputs, orig_target_sizes)
return outputs
model = Model()
data = torch.rand(32, 3, 640, 640)
size = torch.tensor([[640, 640]])
_ = model(data, size)
dynamic_axes = {
"images": {
0: "N",
},
"orig_target_sizes": {0: "N"},
}
output_file = args.resume.replace(".pth", ".onnx") if args.resume else "model.onnx"
torch.onnx.export(
model,
(data, size),
output_file,
input_names=["images", "orig_target_sizes"],
output_names=["labels", "boxes", "scores"],
dynamic_axes=dynamic_axes,
opset_version=16,
verbose=False,
do_constant_folding=True,
)
if args.check:
import onnx
onnx_model = onnx.load(output_file)
onnx.checker.check_model(onnx_model)
print("Check export onnx model done...")
if args.simplify:
import onnx
import onnxsim
dynamic = True
# input_shapes = {'images': [1, 3, 640, 640], 'orig_target_sizes': [1, 2]} if dynamic else None
input_shapes = {"images": data.shape, "orig_target_sizes": size.shape} if dynamic else None
onnx_model_simplify, check = onnxsim.simplify(output_file, test_input_shapes=input_shapes)
onnx.save(onnx_model_simplify, output_file)
print(f"Simplify onnx model {check}...")
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument(
"--config",
"-c",
default="configs/dfine/dfine_hgnetv2_l_coco.yml",
type=str,
)
parser.add_argument(
"--resume",
"-r",
type=str,
)
parser.add_argument(
"--check",
action="store_true",
default=True,
)
parser.add_argument(
"--simplify",
action="store_true",
default=True,
)
args = parser.parse_args()
main(args)
|