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

Function inference

demo/BERT/inference_varseqlen.py:150–214  ·  view source on GitHub ↗
(features, tokens)

Source from the content-addressed store, hash-verified

148
149
150 def inference(features, tokens):
151 global h_output
152
153 _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name
154 "NetworkOutput",
155 ["start_logits", "end_logits", "feature_index"])
156 networkOutputs = []
157
158 eval_time_elapsed = 0
159 for feature_index, feature in enumerate(features):
160 # Copy inputs
161 B = 1
162 S = np.sum(feature.input_mask)
163 input_ids = feature.input_ids[0:S]
164 segment_ids = feature.segment_ids[0:S]
165 cu_seq_lens = np.array([0, S], dtype=np.int32);
166
167 if context.get_binding_shape(0)[0] != S:
168 context.set_binding_shape(0, (S,))
169 if context.get_binding_shape(1)[0] != S:
170 context.set_binding_shape(1, (S,))
171 if context.get_binding_shape(2)[0] != 2:
172 context.set_binding_shape(2, (2,))
173 if context.get_binding_shape(3)[0] != S:
174 context.set_binding_shape(3, (S,))
175
176 h_input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids.ravel()))
177 h_segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids.ravel()))
178 h_cu_seq_lens = cuda.register_host_memory(np.ascontiguousarray(cu_seq_lens.ravel()))
179
180 eval_start_time = time.time()
181 cuda.memcpy_htod_async(d_inputs[0], h_input_ids, stream)
182 cuda.memcpy_htod_async(d_inputs[1], h_segment_ids, stream)
183 cuda.memcpy_htod_async(d_inputs[2], h_cu_seq_lens, stream)
184
185 # Run inference
186 context.execute_async_v2(bindings=[int(d_inp) for d_inp in d_inputs] + [int(d_output)], stream_handle=stream.handle)
187 # Synchronize the stream
188 stream.synchronize()
189 eval_time_elapsed += (time.time() - eval_start_time)
190
191 # Transfer predictions back from GPU
192 cuda.memcpy_dtoh_async(h_output, d_output, stream)
193 stream.synchronize()
194
195 # Only retrieve and post-process the first batch
196 networkOutputs.append(_NetworkOutput(
197 start_logits = np.array(h_output[0:S]),
198 end_logits = np.array(h_output[S:S*2]),
199 feature_index = feature_index
200 ))
201
202 eval_time_elapsed /= len(features)
203
204 # Total number of n-best predictions to generate in the nbest_predictions.json output file
205 n_best_size = 20
206
207 # The maximum length of an answer that can be generated. This is needed

Callers 1

Calls 2

synchronizeMethod · 0.45
appendMethod · 0.45

Tested by

no test coverage detected