capture stdout/stderr
@@ -8,7 +8,8 @@ from dataclasses import dataclass, asdict
from argparse_dataclass import ArgumentParser
from typing import Literal
from enum import Enum
-
+from contextlib import redirect_stderr, redirect_stdout
+import io
class Status:
STARTING = "STARTING"
@@ -65,12 +66,13 @@ class Args:
fp16_mode: bool = False
int8_mode: bool = False
size: int = 224
- pretrained: bool = True
- save_engine: bool = True
+ pretrained: bool = False
+ save_engine: bool = False
num_warmup: int = 10
num_profile: int = 100
- use_cached_engine: bool = True
+ use_cached_engine: bool = False
log_level: Literal["verbose", "info", "error", "warning", "internal_error"] = "error"
+ use_onnx: bool = False
def id(self):
fn = self.model
@@ -83,19 +85,37 @@ class Args:
fn += "-fp16"
if self.int8_mode:
fn += "-int8"
+ if self.use_onnx:
+ fn += "-onnx"
+
return fn
def engine_filepath(self):
- return os.path.join(self.output_dir, self.id() + ".pth")
+ return os.path.join(self.full_output_dir(), self.id() + ".pth")
def metadata_filepath(self):
- return os.path.join(self.output_dir, self.id() + ".json")
+ return os.path.join(self.full_output_dir(), self.id() + ".json")
def write_output(self, output):
with open(self.metadata_filepath(), 'w') as f:
json.dump(output, f, indent=2)
+ def stderr_filepath(self):
+ return os.path.join(self.full_output_dir(), self.id() + ".stdout")
+
+ def stdout_filepath(self):
+ return os.path.join(self.full_output_dir(), self.id() + ".stdout")
+
+ def full_output_dir(self):
+ return os.path.join(self.output_dir, self.id())
+
def run(self):
+ if not os.path.exists(self.full_output_dir()):
+ os.makedirs(self.full_output_dir())
+ with open(self.stderr_filepath(), 'w') as ferr:
+ with open(self.stdout_filepath(), 'w') as fout:
+ with redirect_stderr(ferr):
+ with redirect_stdout(fout):
with torch.no_grad():
model = timm.create_model(self.model, pretrained=self.pretrained)
model = model.cuda().eval()
@@ -117,10 +137,9 @@ class Args:
[data],
fp16_mode=self.fp16_mode,
int8_mode=self.int8_mode,
- log_level=log_level_to_trt(self.log_level)
+ log_level=log_level_to_trt(self.log_level),
+ use_onnx=self.use_onnx
)
- if not os.path.exists(self.output_dir):
- os.makedirs(self.output_dir)
if self.save_engine:
print(f"Saving engine to {self.engine_filepath()}...")
