Unify output tensor allocation across all callers and devices #321
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -187,18 +187,34 @@ VideoDecoder::VideoStreamDecoderOptions::VideoStreamDecoderOptions( | |
| } | ||
| } | ||
|
|
||
| torch::Tensor VideoDecoder::allocateEmptyHWCTensorForStream( | ||
| int streamIndex, | ||
| std::optional<int> numFrames) { | ||
| auto metadata = containerMetadata_.streams[streamIndex]; | ||
| auto options = streams_[streamIndex].options; | ||
| auto height = options.height.value_or(*metadata.height); | ||
| auto width = options.width.value_or(*metadata.width); | ||
|
|
||
| auto tensorOptions = torch::TensorOptions() | ||
| .dtype(torch::kUInt8) | ||
| .layout(torch::kStrided) | ||
| .device(options.device.type()); | ||
| if (numFrames.has_value()) { | ||
| return torch::empty({numFrames.value(), height, width, 3}, tensorOptions); | ||
| } else { | ||
| return torch::empty({height, width, 3}, tensorOptions); | ||
| } | ||
| } | ||
|
|
||
| VideoDecoder::BatchDecodedOutput VideoDecoder::allocateBatchDecodedOutput( | ||
|
There was a problem hiding this comment. The previous So I had to make the "constructor" a member itself. There was a problem hiding this comment. I understand that but by making it a member we are letting it access a lot more information than it strictly needs. It only needs height, width, device, and not all the other stuff in VideoDecoder. What's worse is it can accidentally change the state of the VideoDecoder because it's not I think it may be better to pass in the height and width after looking it up in the caller. There was a problem hiding this comment. A general principle I like to apply is that classes should know how to create correct objects, and ideally, constructed objects should be valid. That means classes should have constructors, and the construtor is responsible for making the object valid. So I think we should keep the We may need a We can take advantage of tensors have reasonable move-like semantics. |
||
| int streamIndex, | ||
| int64_t numFrames) { | ||
| BatchDecodedOutput output; | ||
| output.frames = allocateEmptyHWCTensorForStream(streamIndex, numFrames); | ||
| output.ptsSeconds = torch::empty({numFrames}, {torch::kFloat64}); | ||
| output.durationSeconds = torch::empty({numFrames}, {torch::kFloat64}); | ||
| return output; | ||
| } | ||
|
|
||
| VideoDecoder::VideoDecoder() {} | ||
|
|
||
|
|
@@ -841,7 +857,7 @@ VideoDecoder::RawDecodedOutput VideoDecoder::getDecodedOutputWithFilter( | |
|
|
||
| VideoDecoder::DecodedOutput VideoDecoder::convertAVFrameToDecodedOutput( | ||
| VideoDecoder::RawDecodedOutput& rawOutput, | ||
| torch::Tensor preAllocatedOutputTensor) { | ||
| // Convert the frame to tensor. | ||
| DecodedOutput output; | ||
| int streamIndex = rawOutput.streamIndex; | ||
|
|
@@ -875,7 +891,7 @@ VideoDecoder::DecodedOutput VideoDecoder::convertAVFrameToDecodedOutput( | |
| } | ||
|
|
||
| // Note [preAllocatedOutputTensor with swscale and filtergraph]: | ||
| // Callers must pass a pre-allocated tensor, where the output frame tensor will | ||
| // be stored. This parameter is honored in any case, but it only leads to a | ||
| // speed-up when swscale is used. With swscale, we can tell ffmpeg to place the | ||
| // decoded frame directly into `preAllocatedtensor.data_ptr()`. We haven't yet | ||
|
|
@@ -886,50 +902,25 @@ VideoDecoder::DecodedOutput VideoDecoder::convertAVFrameToDecodedOutput( | |
| void VideoDecoder::convertAVFrameToDecodedOutputOnCPU( | ||
| VideoDecoder::RawDecodedOutput& rawOutput, | ||
| DecodedOutput& output, | ||
| torch::Tensor preAllocatedOutputTensor) { | ||
| int streamIndex = rawOutput.streamIndex; | ||
| AVFrame* frame = rawOutput.frame.get(); | ||
| auto& streamInfo = streams_[streamIndex]; | ||
| if (output.streamType == AVMEDIA_TYPE_VIDEO) { | ||
| if (streamInfo.colorConversionLibrary == ColorConversionLibrary::SWSCALE) { | ||
| rawOutput.data = preAllocatedOutputTensor.data_ptr<uint8_t>(); | ||
| convertFrameToBufferUsingSwsScale(rawOutput); | ||
| } else if ( | ||
| streamInfo.colorConversionLibrary == | ||
| ColorConversionLibrary::FILTERGRAPH) { | ||
| auto tmpTensor = convertFrameToTensorUsingFilterGraph(streamIndex, frame); | ||
| preAllocatedOutputTensor.copy_(tmpTensor); | ||
| } else { | ||
| throw std::runtime_error( | ||
| "Invalid color conversion library: " + | ||
| std::to_string(static_cast<int>(streamInfo.colorConversionLibrary))); | ||
| } | ||
| output.frame = preAllocatedOutputTensor; | ||
|
|
||
| } else if (output.streamType == AVMEDIA_TYPE_AUDIO) { | ||
| // TODO: https://github.com/pytorch-labs/torchcodec/issues/85 implement | ||
|
|
@@ -971,8 +962,11 @@ VideoDecoder::DecodedOutput VideoDecoder::getFramePlayedAtTimestampNoDemux( | |
| return seconds >= frameStartTime && seconds < frameEndTime; | ||
| }); | ||
| // Convert the frame to tensor. | ||
| auto streamIndex = rawOutput.streamIndex; | ||
| auto preAllocatedOutputTensor = allocateEmptyHWCTensorForStream(streamIndex); | ||
|
There was a problem hiding this comment. In the light of this: #312, it would be more correct to use rawOutput.frame's dimensions, and not the stream's dimensions |
||
| auto output = | ||
| convertAVFrameToDecodedOutput(rawOutput, preAllocatedOutputTensor); | ||
| output.frame = MaybePermuteHWC2CHW(streamIndex, output.frame); | ||
| return output; | ||
| } | ||
|
|
||
|
|
@@ -1009,15 +1003,17 @@ void VideoDecoder::validateFrameIndex( | |
| VideoDecoder::DecodedOutput VideoDecoder::getFrameAtIndex( | ||
| int streamIndex, | ||
| int64_t frameIndex) { | ||
| auto preAllocatedOutputTensor = allocateEmptyHWCTensorForStream(streamIndex); | ||
| auto output = getFrameAtIndexInternal( | ||
| streamIndex, frameIndex, preAllocatedOutputTensor); | ||
| output.frame = MaybePermuteHWC2CHW(streamIndex, output.frame); | ||
| return output; | ||
| } | ||
|
|
||
| VideoDecoder::DecodedOutput VideoDecoder::getFrameAtIndexInternal( | ||
| int streamIndex, | ||
| int64_t frameIndex, | ||
| torch::Tensor preAllocatedOutputTensor) { | ||
| validateUserProvidedStreamIndex(streamIndex); | ||
| validateScannedAllStreams("getFrameAtIndex"); | ||
|
|
||
|
|
@@ -1057,7 +1053,8 @@ VideoDecoder::BatchDecodedOutput VideoDecoder::getFramesAtIndices( | |
| const auto& streamMetadata = containerMetadata_.streams[streamIndex]; | ||
| const auto& stream = streams_[streamIndex]; | ||
| const auto& options = stream.options; | ||
| BatchDecodedOutput output = | ||
| allocateBatchDecodedOutput(streamIndex, frameIndices.size()); | ||
|
|
||
| auto previousIndexInVideo = -1; | ||
| for (auto f = 0; f < frameIndices.size(); ++f) { | ||
|
|
@@ -1149,8 +1146,8 @@ VideoDecoder::BatchDecodedOutput VideoDecoder::getFramesInRange( | |
| step > 0, "Step must be greater than 0; is " + std::to_string(step)); | ||
|
|
||
| int64_t numOutputFrames = std::ceil((stop - start) / double(step)); | ||
| BatchDecodedOutput output = | ||
| allocateBatchDecodedOutput(streamIndex, numOutputFrames); | ||
|
|
||
| for (int64_t i = start, f = 0; i < stop; i += step, ++f) { | ||
| DecodedOutput singleOut = | ||
|
|
@@ -1189,9 +1186,6 @@ VideoDecoder::getFramesPlayedByTimestampInRange( | |
| "; must be less than or equal to " + std::to_string(maxSeconds) + | ||
| ")."); | ||
|
|
||
| // Special case needed to implement a half-open range. At first glance, this | ||
| // may seem unnecessary, as our search for stopFrame can return the end, and | ||
| // we don't include stopFramIndex in our output. However, consider the | ||
|
|
@@ -1210,7 +1204,7 @@ VideoDecoder::getFramesPlayedByTimestampInRange( | |
| // values of the intervals will map to the same frame indices below. Hence, we | ||
| // need this special case below. | ||
| if (startSeconds == stopSeconds) { | ||
| BatchDecodedOutput output = allocateBatchDecodedOutput(streamIndex, 0); | ||
| output.frames = MaybePermuteHWC2CHW(streamIndex, output.frames); | ||
| return output; | ||
| } | ||
|
|
@@ -1226,6 +1220,7 @@ VideoDecoder::getFramesPlayedByTimestampInRange( | |
| // 2. In order to establish if the start of an interval maps to a particular | ||
| // frame, we need to figure out if it is ordered after the frame's pts, but | ||
| // before the next frames's pts. | ||
| const auto& stream = streams_[streamIndex]; | ||
| auto startFrame = std::lower_bound( | ||
| stream.allFrames.begin(), | ||
| stream.allFrames.end(), | ||
|
|
@@ -1245,7 +1240,8 @@ VideoDecoder::getFramesPlayedByTimestampInRange( | |
| int64_t startFrameIndex = startFrame - stream.allFrames.begin(); | ||
| int64_t stopFrameIndex = stopFrame - stream.allFrames.begin(); | ||
| int64_t numFrames = stopFrameIndex - startFrameIndex; | ||
| BatchDecodedOutput output = | ||
| allocateBatchDecodedOutput(streamIndex, numFrames); | ||
| for (int64_t i = startFrameIndex, f = 0; i < stopFrameIndex; ++i, ++f) { | ||
| DecodedOutput singleOut = | ||
| getFrameAtIndexInternal(streamIndex, i, output.frames[f]); | ||
|
|
@@ -1267,13 +1263,17 @@ VideoDecoder::RawDecodedOutput VideoDecoder::getNextRawDecodedOutputNoDemux() { | |
| } | ||
|
|
||
| VideoDecoder::DecodedOutput VideoDecoder::getNextFrameNoDemux() { | ||
| auto rawOutput = getNextRawDecodedOutputNoDemux(); | ||
| auto streamIndex = rawOutput.streamIndex; | ||
| auto preAllocatedOutputTensor = allocateEmptyHWCTensorForStream(streamIndex); | ||
|
|
||
| auto output = | ||
| convertAVFrameToDecodedOutput(rawOutput, preAllocatedOutputTensor); | ||
| output.frame = MaybePermuteHWC2CHW(streamIndex, output.frame); | ||
| return output; | ||
| } | ||
|
|
||
| VideoDecoder::DecodedOutput VideoDecoder::getNextFrameOutputNoDemuxInternal( | ||
| torch::Tensor preAllocatedOutputTensor) { | ||
| auto rawOutput = getNextRawDecodedOutputNoDemux(); | ||
| return convertAVFrameToDecodedOutput(rawOutput, preAllocatedOutputTensor); | ||
| } | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please double check this:
To allocate a tensor, we need its height and width. To get those, we need a streamIndex[1]. To get the streamIndex, we need either:
RawOutput.[1] well, not strictly true, there are some cases where it's not needed, but the whole point of this PR is to unify that logic.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In the most generic case this isn't true because a single stream can have frames with different heights and widths:
#312
But ignoring that fact, to get the frame dimensions before decoding you do need the stream_index