成功的尝试

scripts/export_trt.py

@@ -0,0 +1,87 @@
+"""Export video UNet backbone to ONNX, then convert to TensorRT engine.
+
+Usage:
+    python scripts/export_trt.py \
+        --ckpt ckpts/unifolm_wma_dual.ckpt.prepared.pt \
+        --config configs/inference/world_model_interaction.yaml \
+        --out_dir trt_engines
+"""
+
+import os
+import sys
+import argparse
+
+import torch
+import tensorrt as trt
+from omegaconf import OmegaConf
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+from unifolm_wma.utils.utils import instantiate_from_config
+from unifolm_wma.trt_utils import export_backbone_onnx
+
+
+def load_model(config_path, ckpt_path):
+    if ckpt_path.endswith('.prepared.pt'):
+        model = torch.load(ckpt_path, map_location='cpu')
+    else:
+        config = OmegaConf.load(config_path)
+        model = instantiate_from_config(config.model)
+        state_dict = torch.load(ckpt_path, map_location='cpu')
+        if 'state_dict' in state_dict:
+            state_dict = state_dict['state_dict']
+        model.load_state_dict(state_dict, strict=False)
+    model.eval().cuda()
+    return model
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ckpt', required=True)
+    parser.add_argument('--config', default='configs/inference/world_model_interaction.yaml')
+    parser.add_argument('--out_dir', default='trt_engines')
+    parser.add_argument('--context_len', type=int, default=95)
+    parser.add_argument('--fp16', action='store_true', default=True)
+    args = parser.parse_args()
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    onnx_path = os.path.join(args.out_dir, 'video_backbone.onnx')
+    engine_path = os.path.join(args.out_dir, 'video_backbone.engine')
+
+    if os.path.exists(onnx_path):
+        print(f">>> ONNX already exists at {onnx_path}, skipping export.")
+        n_outputs = 10
+    else:
+        print(">>> Loading model ...")
+        model = load_model(args.config, args.ckpt)
+        print(">>> Exporting ONNX ...")
+        with torch.no_grad():
+            n_outputs = export_backbone_onnx(model, onnx_path, context_len=args.context_len)
+        del model
+        torch.cuda.empty_cache()
+
+    print(">>> Converting ONNX -> TensorRT engine ...")
+    logger = trt.Logger(trt.Logger.WARNING)
+    builder = trt.Builder(logger)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+    parser = trt.OnnxParser(network, logger)
+
+    if not parser.parse_from_file(os.path.abspath(onnx_path)):
+        for i in range(parser.num_errors):
+            print(f"    ONNX parse error: {parser.get_error(i)}")
+        raise RuntimeError("ONNX parsing failed")
+
+    config = builder.create_builder_config()
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 16 << 30)
+    if args.fp16:
+        config.set_flag(trt.BuilderFlag.FP16)
+
+    engine_bytes = builder.build_serialized_network(network, config)
+    with open(engine_path, 'wb') as f:
+        f.write(engine_bytes)
+
+    print(f"\n>>> Done! Engine saved to {engine_path}")
+    print(f"    Outputs: 1 y + {n_outputs - 1} hs_a tensors")
+
+
+if __name__ == '__main__':
+    main()