r7vme

#!/usr/bin/env python3                                                  
import graphsurgeon as gs         
import tensorflow as tf   
import tensorrt as trt                                                  
import uff                
                                    
if __name__ == "__main__":                                              
  ### USER DEFINED VARIABLES ###  
  data_type = trt.DataType.HALF                                         
  #data_type = trt.DataType.FLOAT                                       

r7vme

  for node in nodes:
    if "LeakyRelu" in node:
      ns[node]=gs.create_plugin_node(name=node,op="LReLU_TRT", negSlope=0.1)
    if "orientation/l2_normalize" in node:
      dynamic_graph.remove(node)
  dynamic_graph.collapse_namespaces(ns)

r7vme

  with ops.name_scope(name, "l2_normalize", [x]) as name:
    x = ops.convert_to_tensor(x, name="x")
    square_sum = math_ops.reduce_sum(math_ops.square(x), axis, keepdims=True)
    x_inv_norm = math_ops.rsqrt(math_ops.maximum(square_sum, epsilon))
    return math_ops.multiply(x, x_inv_norm, name=name)

r7vme

__global__ void sqrtKernel(
    const int n,
    const float* x,
    float* y)
{
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < n; i += gridDim.x * blockDim.x)
    {
        y[i] = sqrtf(x[i]);
    }

r7vme

__global__ void rsqrtKernel(
    const int n,
    const float* x,
    float* y)
{
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < n; i += gridDim.x * blockDim.x)
    {
        y[i] = rsqrtf(x[i]);
    }

r7vme

__global__ void maxKernel(
    const int n,
    const float eps,
    const float* x,
    float* y)
{
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < n; i += gridDim.x * blockDim.x)
    {
        y[i] = fmaxf(x[i], eps);

r7vme

L2NormHelper::L2NormHelper(const void* buffer, size_t length)
{
    const char *d = reinterpret_cast<const char*>(buffer), *a = d;
    op_type  = read<int>(d);
    eps = read<float>(d);
    C = read<int>(d);
    H = read<int>(d);
    W = read<int>(d);
    ASSERT(d == a + length);
}

r7vme

int L2NormHelper::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream)
{
    const void* inputData = inputs[0];
    void* outputData = outputs[0];
    bool status = executeInference(stream, op_type, eps, batchSize, C, H, W, inputData, outputData);
    ASSERT(status == 0);
    return 0;
}

r7vme

IPluginV2* L2NormHelperPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
{
    const PluginField* fields = fc->fields;
    for (int i = 0; i < fc->nbFields; ++i)
    {
        const char* attrName = fields[i].name;
        if (!strcmp(attrName, "op_type"))
        {
            ASSERT(fields[i].type == PluginFieldType::kINT32);
            mOpType = static_cast<int>(*(static_cast<const int*>(fields[i].data)));

r7vme

REGISTER_TENSORRT_PLUGIN(L2NormHelperPluginCreator);