Calibration of Slowfast architecture on UCF101 dataset

calibrator.py

import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from PIL import Image
import ctypes
import tensorrt as trt

from slowfast.config.defaults import get_cfg
from slowfast.datasets import loader

import os

cfg = get_cfg()

cfg.merge_from_file("/workspace/tensorrt/samples/python/introductory_parser_samples/SlowFast/configs/SLOWFAST_8x8_R50-UCF101.yaml")
cfg.NUM_GPUS = 1
cfg.TRAIN.BATCH_SIZE = 1

CHANNEL = 3
HEIGHT = 224
WIDTH = 224
TEMPORAL_FAST = 8
TEMPORAL_SLOW = 32 

class PythonEntropyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, stream):
        trt.IInt8EntropyCalibrator2.__init__(self)       
        self.stream = stream
        self.d_input_fast = cuda.mem_alloc(self.stream.calibration_data_fast.nbytes*2)
        self.d_input_slow = cuda.mem_alloc(self.stream.calibration_data_slow.nbytes*2)
        self.cache_file = "slow_fast.cache"
        self.current_index = 0

    def get_batch_size(self):
        return cfg.TRAIN.BATCH_SIZE

    def get_batch(self, names):
        self.current_index += 1

        batch, labels, _, meta = self.stream.next_batch()
        
        [fast_batch, slow_batch] = batch

        print(slow_batch.numpy().shape)
        print(fast_batch.numpy().shape)
        
        cuda.memcpy_htod(self.d_input_fast, fast_batch.numpy())
        cuda.memcpy_htod(self.d_input_slow, slow_batch.numpy())
        

        print("names: ", names)

        if self.current_index >= 5:
            return None

        return [self.d_input_fast, self.d_input_slow]

    def read_calibration_cache(self):
        # If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None.
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()

    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)


class ImageBatchStream():
    def __init__(self):

        self.loader = loader.construct_loader(cfg, 'val')

        self.calibration_data_fast = np.zeros((cfg.TRAIN.BATCH_SIZE, CHANNEL, TEMPORAL_FAST, HEIGHT, WIDTH), \
                                            dtype=np.float32)

        self.calibration_data_slow = np.zeros((cfg.TRAIN.BATCH_SIZE, CHANNEL,TEMPORAL_SLOW, HEIGHT, WIDTH), \
                                            dtype=np.float32)
        self.batch = 0
            
    def reset(self):
        self.batch = 0
        
    def next_batch(self):
        return iter(self.loader).next()