Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add MergeTaskOutputGPU which keeps GPU data on the device. #3067

Merged
merged 2 commits into from
Mar 4, 2019
Merged

Add MergeTaskOutputGPU which keeps GPU data on the device. #3067

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
62fcd1c
Add MergeTaskOutputGPU which keeps GPU data on the device.
luitjens Mar 2, 2019
fe3f371
Renamed MergeTaskOutputGPU to MergeTaskOutput
luitjens Mar 2, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions src/nnet3/nnet-batch-compute.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -863,6 +863,40 @@ void MergeTaskOutput(
}
KALDI_ASSERT(cur_output_frame == num_output_frames);
}
void MergeTaskOutputGPU(
const std::vector<NnetInferenceTask> &tasks,
CuMatrix<BaseFloat> *output) {
int32 num_tasks = tasks.size(),
num_output_frames = 0,
output_dim = -1;
for (int32 i = 0; i < num_tasks; i++) {
const NnetInferenceTask &task = tasks[i];
num_output_frames += task.num_used_output_frames;
if (i == 0) {
output_dim = (task.output_to_cpu ?
task.output_cpu.NumCols() :
task.output.NumCols());
}
}
KALDI_ASSERT(num_output_frames != 0 && output_dim != 0);
int32 cur_output_frame = 0;
output->Resize(num_output_frames, output_dim);
for (int32 i = 0; i < num_tasks; i++) {
const NnetInferenceTask &task = tasks[i];
int32 skip = task.num_initial_unused_output_frames,
num_used = task.num_used_output_frames;
KALDI_ASSERT(cur_output_frame == task.first_used_output_frame_index);
if (task.output_to_cpu) {
output->RowRange(cur_output_frame, num_used).CopyFromMat(
task.output_cpu.RowRange(skip, num_used));
} else {
output->RowRange(cur_output_frame, num_used).CopyFromMat(
task.output.RowRange(skip, num_used));
}
cur_output_frame += num_used;
}
KALDI_ASSERT(cur_output_frame == num_output_frames);
}


NnetBatchInference::NnetBatchInference(
Expand Down
3 changes: 3 additions & 0 deletions src/nnet3/nnet-batch-compute.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,9 @@ struct NnetBatchComputerOptions: public NnetSimpleComputationOptions {
void MergeTaskOutput(
const std::vector<NnetInferenceTask> &tasks,
Matrix<BaseFloat> *output);
void MergeTaskOutputGPU(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. You can just call it MergeTaskOutput() and let the function overloading take care of it, that's what I normally do for CPU/GPU things. Incidentically, if it had the suffix, it would be MergeTaskOutputGpu()... it's a Google style guide thing.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed

const std::vector<NnetInferenceTask> &tasks,
CuMatrix<BaseFloat> *output);

/**
This class does neural net inference in a way that is optimized for GPU use:
Expand Down