Random sampling while reading tfrecord

TL;DR
I have 1000x3 float data in tfrecord and I want to read it as 500x3 data where 500 is randomly sampled from 1000. I want it to be random every time I read it.

Hi, I'm using PyTorch + DALI + tfrecord and GPUs.

I stored 3D point cloud (xyz) from several frames in the tfrecord file.
The number of points in each frame is different.
For example, at frame 1, I have 1000 points so the information is stored as 1000x3 in float type.
At frame 2, I have 2000 points so it is 2000x3 data.
Therefore, I use
"cloud" : tfrec.VarLenFeature(tfrec.float32, 0.)
to load the point cloud in DALI's ops.TFRecordReader.

To get this data into PyTorch's tensor, I first pad constant values using ops.Pad to match the maximum number of points in the batch.
In the above example, the second frame has the biggest number of points (2000).
Therefore, I pad the first point cloud with constant 1000x3 values.
A batch in PyTorch tensor with batch_size=2 will look like this: 2 x 2000 x 3.

For each iteration, I want to randomly sample 500 points at each frame to make 2 x 500 x 3 as an input to the network.
In order to do that, I should remember the number of points at each frame.
For example, for frame 1, although the batch is loaded as 2000x3 shape, only first 1000x3 are real points and the rest of 1000x3 are padded constant.
Therefore, I sample 500 points from 0~999 range.

This approach works, however, it seems it has an overhead.
The GPU usage is around 60% on the nvidia-smi.

I think it would be great if we can directly do random sampling when we load the data from tfrecord.
I checked two functions in DALI.
One is Slice and the other is ElementExtract.
For the Slice operator, it randomly gets a chunk. It does not look like a random access that I described above.
For ElementExtract, I'm not sure whether it can achieve my goal.

How can I do this?

Start by reviewing the DALI ops.TFRecordReader, ops.Pad, ops.Slice, and ops.ElementExtract entry points and the current PyTorch pipeline described in the issue. Determine whether variable-length records can support per-frame random sampling during loading; done means producing a fixed 500x3 sample from each frame on every read without sampling padded values.