@@ -113,10 +113,6 @@ void TrtMTR::initCudaPtr(const AgentData & agent_data, const PolylineData & poly
113
113
d_out_trajectory_ =
114
114
cuda::make_unique<float []>(num_target_ * num_mode_ * num_future_ * PredictedStateDim);
115
115
116
- h_out_score_ = std::make_unique<float []>(sizeof (float ) * num_target_ * num_mode_);
117
- h_out_trajectory_ = std::make_unique<float []>(
118
- sizeof (float ) * num_target_ * num_mode_ * num_future_ * PredictedStateDim);
119
-
120
116
if (builder_->isDynamic ()) {
121
117
// trajectory: (B, N, T, Da)
122
118
builder_->setBindingDimensions (
@@ -198,20 +194,27 @@ bool TrtMTR::postProcess(
198
194
num_target_, num_mode_, num_future_, agent_data.state_dim (), d_target_state_.get (),
199
195
PredictedStateDim, d_out_trajectory_.get (), stream_));
200
196
197
+ h_out_score_.clear ();
198
+ h_out_trajectory_.clear ();
199
+ h_out_score_.reserve (num_target_ * num_mode_);
200
+ h_out_trajectory_.reserve (num_target_ * num_mode_ * num_future_ * PredictedStateDim);
201
+
201
202
CHECK_CUDA_ERROR (cudaMemcpyAsync (
202
- h_out_score_.get (), d_out_score_.get (), sizeof (float ) * num_target_ * num_mode_,
203
+ h_out_score_.data (), d_out_score_.get (), sizeof (float ) * num_target_ * num_mode_,
203
204
cudaMemcpyDeviceToHost, stream_));
204
205
CHECK_CUDA_ERROR (cudaMemcpyAsync (
205
- h_out_trajectory_.get (), d_out_trajectory_.get (),
206
+ h_out_trajectory_.data (), d_out_trajectory_.get (),
206
207
sizeof (float ) * num_target_ * num_mode_ * num_future_ * PredictedStateDim,
207
208
cudaMemcpyDeviceToHost, stream_));
208
209
209
210
trajectories.reserve (num_target_);
210
211
for (auto b = 0 ; b < num_target_; ++b) {
211
- const auto score_ptr = h_out_score_.get () + b * num_mode_;
212
- const auto trajectory_ptr =
213
- h_out_trajectory_.get () + b * num_mode_ * num_future_ * PredictedStateDim;
214
- trajectories.emplace_back (score_ptr, trajectory_ptr, num_mode_, num_future_);
212
+ const auto score_itr = h_out_score_.cbegin () + b * num_mode_;
213
+ std::vector<float > scores (score_itr, score_itr + num_mode_);
214
+ const auto mode_itr =
215
+ h_out_trajectory_.cbegin () + b * num_mode_ * num_future_ * PredictedStateDim;
216
+ std::vector<float > modes (mode_itr, mode_itr + num_mode_ * num_future_ * PredictedStateDim);
217
+ trajectories.emplace_back (scores, modes, num_mode_, num_future_);
215
218
}
216
219
return true ;
217
220
}
0 commit comments