Created
July 4, 2020 12:54
-
-
Save kunalgoyal9/1cc1079693786ec633496b8d3e810bd6 to your computer and use it in GitHub Desktop.
Calibration of Slowfast architecture on UCF101 dataset
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
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment