MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / inference

Function inference

demo/BERT/inference_c.py:132–181  ·  view source on GitHub ↗
(features, tokens)

Source from the content-addressed store, hash-verified

130 bert.prepare(1)
131
132 def inference(features, tokens):
133
134 _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name
135 "NetworkOutput",
136 ["start_logits", "end_logits", "feature_index"])
137 networkOutputs = []
138
139 eval_time_elapsed = 0
140 for feature_index, feature in enumerate(features):
141 # Copy inputs
142 input_ids = np.ascontiguousarray(feature.input_ids.ravel())
143 segment_ids = np.ascontiguousarray(feature.segment_ids.ravel())
144 input_mask = np.ascontiguousarray(feature.input_mask.ravel())
145
146 eval_start_time = time.time()
147
148 # Run inference
149 h_output = bert.run(input_ids, segment_ids, input_mask)
150 eval_time_elapsed += (time.time() - eval_start_time)
151
152
153 # Data Post-processing
154 if len(h_output.shape) == 1:
155 S = int(h_output.shape[0] / 2)
156 networkOutputs.append(_NetworkOutput(
157 start_logits = np.array(h_output[0:S]),
158 end_logits = np.array(h_output[S:S*2]),
159 feature_index = feature_index
160 ))
161 else:
162 for index, batch in enumerate(h_output):
163 networkOutputs.append(_NetworkOutput(
164 start_logits = np.array(batch.squeeze()[:, 0]),
165 end_logits = np.array(batch.squeeze()[:, 1]),
166 feature_index = feature_index
167 ))
168
169 eval_time_elapsed /= len(features)
170
171 # Total number of n-best predictions to generate in the nbest_predictions.json output file
172 n_best_size = 20
173
174 # The maximum length of an answer that can be generated. This is needed
175 # because the start and end predictions are not conditioned on one another
176 max_answer_length = 30
177
178 prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,
179 networkOutputs, args.n_best_size, args.max_answer_length)
180
181 return eval_time_elapsed, prediction, nbest_json
182
183 def print_single_query(eval_time_elapsed, prediction, nbest_json):
184 print("------------------------")

Callers 1

inference_c.pyFile · 0.70

Calls 2

runMethod · 0.45
appendMethod · 0.45

Tested by

no test coverage detected