Sampling Random Frames from Short Video Clips

I am trying to use the VideoLoader to extract individual frames from video clips (~10 sec each). I was previously carrying out this task by extracting frames and saving them before training. I need to be able to augment these frames as well. Is there currently a way to do this?

I used the ElementExtract operator to turn the sequence into a single frame, but the training process was extremely slow (don't have a metric but it took ~3 seconds for a single batch of size 64). In addition, changing the step parameter in the VideoReader changes the epoch_size parameter of my pipeline. Lastly, the iterator has to be reset after ~12 iterations for some reason, so each epoch only runs 12 batches. Not sure what's going on, but it would be great to get some guidance on what might be happening here. Attaching my code as well.

import os
import numpy as np

from nvidia.dali.pipeline import Pipeline
import nvidia.dali.ops as ops
import nvidia.dali.types as types
from nvidia.dali.plugin.pytorch import DALIGenericIterator, TorchPythonFunction

def reshape(input_seq):
    return input_seq.view(-1, input_seq.size()[-3], input_seq.size()[-2], input_seq.size()[-1])

class Training_Pipeline(Pipeline):
    def __init__(self, num_threads, device_id, annotation_path, shuffle, sequence_length=1, batch_size=64):
        super(Training_Pipeline, self).__init__(batch_size, num_threads, device_id)
        self.sequence_length = sequence_length

        #TODO figure out these parameters
        self.input = ops.VideoReader(device='gpu', file_list=annotation_path, step = 10, stride=10,
                sequence_length=self.sequence_length, image_type=types.RGB, shard_id=0,
                num_shards=1, random_shuffle=shuffle, initial_fill=1, normalized=True)

        self.small_rotation = ops.Uniform(range=(-10.0, 10.0))
        self.brightness_change = ops.Uniform(range=(0.8, 1.2))
        
        self.rotate = ops.Rotate(device = 'gpu')
        self.brightness = ops.Brightness(device = 'gpu')
        self.resize = ops.Resize(device = "gpu", image_type = types.RGB,
                interp_type = types.INTERP_LINEAR, resize_x = 224, resize_y = 224)
        self.element_extract = ops.ElementExtract(element_map=0, device='gpu')

        #self.reshape = TorchPythonFunction(reshape, device='cpu')

    def transform(self, inputs):
        # remove arbitrary 2nd dimension (we are only selecting 1 frame per sequence)
        images = self.element_extract(inputs)
        images = self.resize(images)
        if np.random.randint(0, 2):
            images = ops.Flip(device='gpu', horizontal = 0, vertical = 1)(images)

        if np.random.randint(0, 2):
            images = ops.Flip(device='gpu', horizontal = 1, vertical = 0)(images)

        images = ops.Rotate(device = 'gpu', angle = 90 * np.random.randint(0, 4))(images)
        images = self.rotate(images, angle = self.small_rotation())
        images = self.brightness(images, brightness = self.brightness_change())

        return images

    def define_graph(self):
        images, labels = self.input(name="Reader")
        images = self.transform(images)

        return images, labels

class DALILoader():
    def __init__(self, num_threads, device_id, batch_size, annotation_path, shuffle):
        self.pipeline = Training_Pipeline(batch_size=batch_size, num_threads=num_threads,
                device_id=device_id, annotation_path=annotation_path, shuffle=shuffle)

        self.pipeline.build()
        self.epoch_size = self.pipeline.epoch_size("Reader")
        self.dali_iterator = DALIGenericIterator(self.pipeline, ["data", "label"], self.epoch_size,
             auto_reset=True)

    def __len__(self):
        return int(self.epoch_size)

    def __iter__(self):
        return self.dali_iterator.__iter__()

# wrap dataset in dataloader
def prepare_dataset(args):
    #TODO figure out num_threads parameter
    if args.cuda:
        kwargs = {'num_threads': 16, 'device_id': 0, 'annotation_path': os.path.join(args.annotation_path,
            "train_annots.txt"), 'shuffle': True, 'batch_size': args.batch_size}
        val_kwargs = {'num_threads': 16, 'device_id': 0, 'annotation_path': args.test_annotation_path,
                'shuffle': True, 'batch_size': args.val_batch_size}
        test_kwargs = {'num_threads': 8, 'device_id': 0, 'annotation_path': args.test_annotation_path,
                'shuffle': True, 'batch_size': args.test_batch_size}

    else:
        kwargs = {'num_threads': 2, 'device_id': 0, 'annotation_path': args.annotation_path,
                'shuffle': True, 'batch_size': args.batch_size}
        val_kwargs = {'num_threads': 8, 'device_id': 0, 'annotation_path': args.test_annotation_path,
                'shuffle': True, 'batch_size': args.val_batch_size}
        test_kwargs = {'num_threads': 8, 'device_id': 0, 'annotation_path': args.test_annotation_path,
                'shuffle': True, 'batch_size': args.test_batch_size}

    train_loader, val_loader, test_loaders = None, None, None

    if not args.test_only:
        #train_pipe = Training_Pipeline(**kwargs)
        #train_pipe.build()
        train_loader = DALILoader(**kwargs)
        #train_loader = DALIGenericIterator(train_pipe, ['data', 'label'], train_pipe.epoch_size("Reader"))

    #val_dataset = DatasetClass(args.val_data_dir, datasets = args.datasets, categories = args.categories,
    #        extra_channel = args.extra_channel, transform = transforms.Compose(resize_normalize_only))

    #val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.test_batch_size, shuffle=False, **test_kwargs)

    #test_dataset = DatasetClass(args.test_data_dir, datasets = args.datasets,
    #        categories = args.categories, extra_channel = args.extra_channel,
    #        transform = transforms.Compose(resize_normalize_only))

    #if args.tta:
    #    test_dataset_augmented = DatasetClass(args.test_data_dir, datasets = args.datasets,
    #            categories = args.categories, extra_channel = args.extra_channel,
    #            transform = transforms.Compose(test_composed))

    #    test_dataset2_augmented = DatasetClass(args.test_data_dir, datasets = args.datasets,
    #            categories = args.categories, extra_channel = args.extra_channel,
    #            transform = transforms.Compose(test_composed))

    #    test_loaders = [torch.utils.data.DataLoader(test_dataset, batch_size=args.test_batch_size,
    #            shuffle=False, **test_kwargs),
    #            torch.utils.data.DataLoader(test_dataset_augmented, batch_size=args.test_batch_size,
    #                shuffle=False, **test_kwargs),
    #            torch.utils.data.DataLoader(test_dataset2_augmented, batch_size=args.test_batch_size,
    #                shuffle=False, **test_kwargs)]

    #else:
    #    test_loaders = [torch.utils.data.DataLoader(test_dataset, batch_size=args.test_batch_size,
    #            shuffle=False, **test_kwargs)]

    return train_loader, val_loader, test_loaders

Start with the Training_Pipeline.define_graph and DALILoader code shown in the issue, focusing on VideoReader, ElementExtract, and DALIGenericIterator behavior. Reproduce the short-clip sampling, epoch-size, throughput, and iterator-reset observations; done means documenting or implementing a supported approach for random frame sampling with augmentation and correct epoch iteration.