Commit e9667fa0 authored by Glydzo's avatar Glydzo
Browse files

First working example

parent 9c647030
src/__pycache__
# CNN-with-TRT
# Requirements
Cuda 11.0 update 1 : [link](https://developer.nvidia.com/cuda-11.0-update1-download-archive)
-> Use the local runfile install (you have to previously disable Nouveau driver and unistall all Nvidia drivers).
TensorRT 8.0 GA : [link](https://developer.nvidia.com/nvidia-tensorrt-8x-download)
-> Download TensorRT 8.0.1 GA for Linux x86_64 and CUDA 11.3 TAR package and follow the [setup guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html#installing-tar).
# Python setting up
First, be sure you have Python 3.6 installed (if you are using ubuntu 18.04, it is native).
'''
python3 -m venv venv
source venv/bin/activate
python -m pip install -U pip
python -m pip install tensorrt-*-cp3x-none-linux_x86_64.whl
python -m pip install uff-0.6.9-py2.py3-none-any.whl
python -m pip install graphsurgeon-0.4.5-py2.py3-none-any.whl
python -m pip install onnx_graphsurgeon-0.2.6-py2.py3-none-any.whl
python -m pip install tensorflow==2.4.0
python -m pip install tqdm
'''
# Useful links
[Leveraging TensorFlow-TensorRT integration for Low latency Inference](https://blog.tensorflow.org/2021/01/leveraging-tensorflow-tensorrt-integration.html)
......
model_checkpoint_path: "my_weights"
all_model_checkpoint_paths: "my_weights"
import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Conv2D, MaxPooling2D, Flatten, Activation, add, Concatenate, BatchNormalization
from tensorflow.keras.models import Model
def conv_block(input_tensor,
kernel_size,
nb_filter,
name
):
x = Conv2D(nb_filter, (kernel_size, kernel_size), padding="same", activation="relu", name=name + "_conv1")(
input_tensor)
x = Conv2D(nb_filter, (kernel_size, kernel_size), padding="same", name=name + "_conv2")(x)
# Compute one convolution to resize correctly the input_tensor as to compute the add operation
input_tensor = Conv2D(nb_filter, (1, 1), padding="same")(input_tensor)
x = add([x, input_tensor])
x1 = Activation('relu')(x)
x = Conv2D(nb_filter, (kernel_size, kernel_size), padding="same", activation="relu", name=name + "_conv3")(x1)
x = Conv2D(nb_filter, (kernel_size, kernel_size), padding="same", name=name + "_conv4")(x)
x = add([x, x1])
x = Activation('relu')(x)
return x
def model_technicolor_vector_multi_qp():
img_input = Input(shape=(68, 68, 1))
qp_input = Input(shape=(1,), name='qp_input')
x = Conv2D(16, (3, 3), padding="same", activation= "relu", name='init_conv')(img_input)
x = MaxPooling2D((2, 2), padding="same", strides=(2, 2))(x)
x = conv_block(x, 3, 24, name='block1')
x = MaxPooling2D((3, 3), strides=(3, 3), padding="same")(x)
x = conv_block(x, 3, 32, name='block2')
x = MaxPooling2D((3, 3), strides=(3, 3), padding="same")(x)
x = conv_block(x, 3, 48, name='block3')
x = MaxPooling2D((3, 3), strides=(3, 3), padding="same")(x)
x = Flatten()(x)
x = Concatenate()([x, qp_input])
output_luma = Dense(480, activation='sigmoid', name='fc480')(x)
# Add QP to X here before creating the model
model = Model(inputs = [img_input, qp_input], outputs = output_luma, name='archiGlobal')
return model
weights_path = '../res/weights/'
dataset_path = '../res/dataset/'
import time
import os
# 0 = all messages are logged (default behavior)
# 1 = INFO messages are not printed
# 2 = INFO and WARNING messages are not printed
# 3 = INFO, WARNING, and ERROR messages are not printed
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
import tensorflow as tf
from CNN import model_technicolor_vector_multi_qp
import numpy as np
from tqdm import tqdm
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
except RuntimeError as e:
print(e)
with tf.device('/GPU:0'):
model = model_technicolor_vector_multi_qp()
# Either load the weights from the h5 file, or load them from the native format using the expect_partial() method of load_weights(...)
model.load_weights(weights_path + 'weights.h5')
#model.load_weights(weights_path + 'my_weights').expect_partial()
#model.summary()
image = np.array([np.load('../res/dataset/1000008.npy')])
qp = np.array([22])
elapsed_times = []
for i in tqdm(range(10000)):
start_time = time.time()
prediction = model.predict([image,qp])
end_time = time.time()
elapsed_times.append((end_time - start_time) * 1000)
print('Inferences finished ! Dimensions : ', prediction.shape, ' / Average execution time : %.3f ms ' % (sum(elapsed_times)/len(elapsed_times)), sep="")
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment