(
self,
model,
exp,
trt_file=None,
decoder=None,
device=torch.device("cpu"),
fp16=False
)
| 116 | |
| 117 | class Predictor(object): |
| 118 | def __init__( |
| 119 | self, |
| 120 | model, |
| 121 | exp, |
| 122 | trt_file=None, |
| 123 | decoder=None, |
| 124 | device=torch.device("cpu"), |
| 125 | fp16=False |
| 126 | ): |
| 127 | self.model = model |
| 128 | self.decoder = decoder |
| 129 | self.num_classes = exp.num_classes |
| 130 | self.confthre = exp.test_conf |
| 131 | self.nmsthre = exp.nmsthre |
| 132 | self.test_size = exp.test_size |
| 133 | self.device = device |
| 134 | self.fp16 = fp16 |
| 135 | if trt_file is not None: |
| 136 | from torch2trt import TRTModule |
| 137 | |
| 138 | model_trt = TRTModule() |
| 139 | model_trt.load_state_dict(torch.load(trt_file)) |
| 140 | |
| 141 | x = torch.ones((1, 3, exp.test_size[0], exp.test_size[1]), device=device) |
| 142 | self.model(x) |
| 143 | self.model = model_trt |
| 144 | self.rgb_means = (0.485, 0.456, 0.406) |
| 145 | self.std = (0.229, 0.224, 0.225) |
| 146 | |
| 147 | def inference(self, img, timer): |
| 148 | img_info = {"id": 0} |
nothing calls this directly
no outgoing calls
no test coverage detected