Add: Wideresnet (#4) (#518)

* Adding WideResnet C++ (#2) * initialize wideresnet50 * add: wideresnet50 c++ code * add: wide resnet python (#3) * add: wide resnet python * fix: typo Co-authored-by: makaveli <39617050+makaveli10@users.noreply.github.com> Co-authored-by: makaveli <39617050+makaveli10@users.noreply.github.com>
2021-04-30 08:00:40 +05:30 · 2021-04-30 08:00:40 +05:30 · 4ffc56a99b
commit 4ffc56a99b
parent 1f7672ce5a
4 changed files with 675 additions and 5 deletions
--- a/resnet/CMakeLists.txt
+++ b/resnet/CMakeLists.txt
@ -29,5 +29,9 @@ add_executable(resnext50 ${PROJECT_SOURCE_DIR}/resnext50_32x4d.cpp)
 target_link_libraries(resnext50 nvinfer)
 target_link_libraries(resnext50 cudart)

+add_executable(wideresnet50 ${PROJECT_SOURCE_DIR}/wideresnet50.cpp)
+target_link_libraries(wideresnet50 nvinfer)
+target_link_libraries(wideresnet50 cudart)
+
 add_definitions(-O2 -pthread)

--- a/resnet/README.md
+++ b/resnet/README.md
@ -4,6 +4,8 @@ ResNet-18 and ResNet-50 model from "Deep Residual Learning for Image Recognition

 For the Pytorch implementation, you can refer to [pytorchx/resnet](https://github.com/wang-xinyu/pytorchx/tree/master/resnet)

+Wide Resnet-50 model from "Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf> . For the Pytorch implementation, you can refer to [BlueMirrors/torchtrtz](https://github.com/BlueMirrors/torchtrtz)
+
 Following tricks are used in this resnet, nothing special, residual connection and batchnorm are used.

 - Batchnorm layer, implemented with scale layer.
@ -11,7 +13,9 @@ Following tricks are used in this resnet, nothing special, residual connection a
 ## TensorRT C++ API

 ```
-// 1. generate resnet18.wts or resnet50.wts from [pytorchx/resnet](https://github.com/wang-xinyu/pytorchx/tree/master/resnet)
+// 1a. generate resnet18.wts or resnet50.wts from [pytorchx/resnet](https://github.com/wang-xinyu/pytorchx/tree/master/resnet)
+
+// 1b. generate wide_resnet50.wts from [BlueMirrors/torchtrtz](https://github.com/BlueMirrors/torchtrtz)

 // 2. put resnet18.wts or resnet50.wts into tensorrtx/resnet

@ -35,16 +39,29 @@ or
 sudo ./resnet50 -s   // serialize model to plan file i.e. 'resnet50.engine'
 sudo ./resnet50 -d   // deserialize plan file and run inference

+or

-// 4. see if the output is same as pytorchx/resnet
+sudo ./resnext50 -s   // serialize model to plan file i.e. 'resnext50.engine'
+sudo ./resnext50 -d   // deserialize plan file and run inference
+
+or
+
+sudo ./wide_resnet50 -s   // serialize model to plan file i.e. 'wide_resnet50.engine'
+sudo ./wide_resnet50 -d   // deserialize plan file and run inference
+
+
+// 4. see if the output is same as 
+- [pytorchx/resnet](https://github.com/wang-xinyu/pytorchx/tree/master/resnet) - for resnet18, resnet50, resnext50
+- [BlueMirrors/torchtrtz](https://github.com/BlueMirrors/torchtrtz) - for wide_resnet50
 ```

 ### TensorRT Python API

 ```
-# 1. generate resnet50.wts from [pytorchx/resnet](https://github.com/wang-xinyu/pytorchx/tree/master/resnet)
+# 1a. generate resnet50.wts from [pytorchx/resnet](https://github.com/wang-xinyu/pytorchx/tree/master/resnet)
+# 1b. generate wide_resnet50.wts from [BlueMirrors/torchtrtz](https://github.com/BlueMirrors/torchtrtz)

-# 2. put resnet50.wts into tensorrtx/resnet
+# 2. put resnet50.wts or wide_resnet50.wts into tensorrtx/resnet

 # 3. install Python dependencies (tensorrt/pycuda/numpy)

@ -53,5 +70,12 @@ cd tensorrtx/resnet
 python resnet50.py -s   // serialize model to plan file i.e. 'resnet50.engine'
 python resnet50.py -d   // deserialize plan file and run inference

-# 4. see if the output is same as pytorchx/resnet
+or 
+
+python wide_resnet50.py -s   // serialize model to plan file i.e. 'wide_resnet50.engine'
+python wide_resnet50.py -d   // deserialize plan file and run inference
+
+# 4. see if the output is same as 
+- pytorchx/resnet - for resnet50
+- BlueMirrors/torchtrtz - for wide_resnet50
 ```
--- a/resnet/wide_resnet50.py
+++ b/resnet/wide_resnet50.py
@ -0,0 +1,275 @@
+import os
+import sys
+import struct
+import argparse
+
+import numpy as np
+import pycuda.autoinit
+import pycuda.driver as cuda
+import tensorrt as trt
+
+BATCH_SIZE = 1
+INPUT_H = 224
+INPUT_W = 224
+OUTPUT_SIZE = 1000
+BS = 1
+INPUT_BLOB_NAME = "data"
+OUTPUT_BLOB_NAME = "prob"
+EPS = 1e-5
+
+WEIGHT_PATH = "./wide_resnet50.wts"
+ENGINE_PATH = "./wide_resnet50.engine"
+
+TRT_LOGGER = trt.Logger(trt.Logger.INFO)
+
+
+def load_weights(file):
+    print(f"Loading weights: {file}")
+
+    assert os.path.exists(file), 'Unable to load weight file.'
+
+    weight_map = {}
+    with open(file, "r") as f:
+        lines = [line.strip() for line in f]
+    count = int(lines[0])
+    assert count == len(lines) - 1
+    for i in range(1, count + 1):
+        splits = lines[i].split(" ")
+        name = splits[0]
+        cur_count = int(splits[1])
+        assert cur_count + 2 == len(splits)
+        values = []
+        for j in range(2, len(splits)):
+            # hex string to bytes to float
+            values.append(struct.unpack(">f", bytes.fromhex(splits[j])))
+        weight_map[name] = np.array(values, dtype=np.float32)
+
+    return weight_map
+
+
+def addBatchNorm2d(network, weight_map, inputs, layer_name, eps):
+    gamma = weight_map[layer_name + ".weight"]
+    beta = weight_map[layer_name + ".bias"]
+    mean = weight_map[layer_name + ".running_mean"]
+    var = weight_map[layer_name + ".running_var"]
+    print(layer_name + " " +  str(len(weight_map[layer_name + ".running_var"])))
+    var = np.sqrt(var + eps)
+
+    scale = gamma / var
+    shift = -mean / var * gamma + beta
+    return network.add_scale(input=inputs,
+                             mode=trt.ScaleMode.CHANNEL,
+                             shift=shift,
+                             scale=scale)
+
+
+def bottleneck(network, weight_map, input, in_channels, out_channels, stride, layer_name):
+    # empty weights for bias
+    emptywts = trt.Weights()
+
+    conv1 = network.add_convolution(input=input,
+                                    num_output_maps=out_channels,
+                                    kernel_shape=(1, 1),
+                                    kernel=weight_map[layer_name + "conv1.weight"],
+                                    bias=emptywts)
+    assert conv1
+
+    bn1 = addBatchNorm2d(network, weight_map, conv1.get_output(0), layer_name + "bn1", EPS)
+    assert bn1
+
+    relu1 = network.add_activation(bn1.get_output(0), type=trt.ActivationType.RELU)
+    assert relu1
+
+    conv2 = network.add_convolution(input=relu1.get_output(0),
+                                    num_output_maps=out_channels,
+                                    kernel_shape=(3, 3),
+                                    kernel=weight_map[layer_name + "conv2.weight"],
+                                    bias=emptywts)
+    assert conv2
+    conv2.stride = (stride, stride)
+    conv2.padding = (1, 1)
+
+    bn2 = addBatchNorm2d(network, weight_map, conv2.get_output(0),
+                         layer_name + "bn2", EPS)
+    assert bn2
+
+    relu2 = network.add_activation(bn2.get_output(0),
+                                   type=trt.ActivationType.RELU)
+    assert relu2
+
+    conv3 = network.add_convolution(input=relu2.get_output(0),
+                                    num_output_maps=out_channels * 2,
+                                    kernel_shape=(1, 1),
+                                    kernel=weight_map[layer_name + "conv3.weight"],
+                                    bias=emptywts)
+    assert conv3
+
+    bn3 = addBatchNorm2d(network, weight_map, conv3.get_output(0), layer_name + "bn3", EPS)
+    assert bn3
+
+    if stride != 1 or in_channels != 2 * out_channels:
+        conv4 = network.add_convolution(
+            input=input,
+            num_output_maps=out_channels * 2,
+            kernel_shape=(1, 1),
+            kernel=weight_map[layer_name + "downsample.0.weight"],
+            bias=emptywts)
+        assert conv4
+        conv4.stride = (stride, stride)
+
+        bn4 = addBatchNorm2d(network, weight_map, conv4.get_output(0), layer_name + "downsample.1", EPS)
+        assert bn4
+
+        ew1 = network.add_elementwise(bn4.get_output(0), bn3.get_output(0),
+                                      trt.ElementWiseOperation.SUM)
+    else:
+        ew1 = network.add_elementwise(input, bn3.get_output(0), trt.ElementWiseOperation.SUM)
+    assert ew1
+
+    relu3 = network.add_activation(ew1.get_output(0), type=trt.ActivationType.RELU)
+    assert relu3
+
+    return relu3
+
+
+def create_engine(maxBatchSize, builder, config, dt):
+    weight_map = load_weights(WEIGHT_PATH)
+    network = builder.create_network()
+
+    data = network.add_input(INPUT_BLOB_NAME, dt, (3, INPUT_H, INPUT_W))
+    assert data
+
+    # empty weights for bias
+    emptywts = trt.Weights()
+
+    conv1 = network.add_convolution(input=data,
+                                    num_output_maps=64,
+                                    kernel_shape=(7, 7),
+                                    kernel=weight_map["conv1.weight"],
+                                    bias=emptywts)
+    assert conv1
+    conv1.stride = (2, 2)
+    conv1.padding = (3, 3)
+
+    bn1 = addBatchNorm2d(network, weight_map, conv1.get_output(0), "bn1", EPS)
+    assert bn1
+
+    relu1 = network.add_activation(bn1.get_output(0), type=trt.ActivationType.RELU)
+    assert relu1
+
+    pool1 = network.add_pooling(input=relu1.get_output(0),
+                                window_size=trt.DimsHW(3, 3),
+                                type=trt.PoolingType.MAX)
+    assert pool1
+    pool1.stride = (2, 2)
+    pool1.padding = (1, 1)
+
+    x = bottleneck(network, weight_map, pool1.get_output(0), 64, 128, 1, "layer1.0.")
+    x = bottleneck(network, weight_map, x.get_output(0), 256, 128, 1, "layer1.1.")
+    x = bottleneck(network, weight_map, x.get_output(0), 256, 128, 1, "layer1.2.")
+
+    x = bottleneck(network, weight_map, x.get_output(0), 256, 256, 2, "layer2.0.")
+    x = bottleneck(network, weight_map, x.get_output(0), 512, 256, 1, "layer2.1.")
+    x = bottleneck(network, weight_map, x.get_output(0), 512, 256, 1, "layer2.2.")
+    x = bottleneck(network, weight_map, x.get_output(0), 512, 256, 1, "layer2.3.")
+
+    x = bottleneck(network, weight_map, x.get_output(0), 512, 512, 2, "layer3.0.")
+    x = bottleneck(network, weight_map, x.get_output(0), 1024, 512, 1, "layer3.1.")
+    x = bottleneck(network, weight_map, x.get_output(0), 1024, 512, 1, "layer3.2.")
+    x = bottleneck(network, weight_map, x.get_output(0), 1024, 512, 1, "layer3.3.")
+    x = bottleneck(network, weight_map, x.get_output(0), 1024, 512, 1, "layer3.4.")
+    x = bottleneck(network, weight_map, x.get_output(0), 1024, 512, 1, "layer3.5.")
+
+    x = bottleneck(network, weight_map, x.get_output(0), 1024, 1024, 2, "layer4.0.")
+    x = bottleneck(network, weight_map, x.get_output(0), 2048, 1024, 1, "layer4.1.")
+    x = bottleneck(network, weight_map, x.get_output(0), 2048, 1024, 1, "layer4.2.")
+
+    pool2 = network.add_pooling(x.get_output(0),
+                                window_size=trt.DimsHW(7, 7),
+                                type=trt.PoolingType.AVERAGE)
+    assert pool2
+    pool2.stride = (1, 1)
+
+    fc1 = network.add_fully_connected(input=pool2.get_output(0),
+                                      num_outputs=OUTPUT_SIZE,
+                                      kernel=weight_map['fc.weight'],
+                                      bias=weight_map['fc.bias'])
+    assert fc1
+
+    fc1.get_output(0).name = OUTPUT_BLOB_NAME
+    network.mark_output(fc1.get_output(0))
+
+    # Build engine
+    builder.max_batch_size = maxBatchSize
+    builder.max_workspace_size = 1 << 20
+    engine = builder.build_engine(network, config)
+    print("build out")
+    del network
+    del weight_map
+
+    return engine
+
+
+def APIToModel(maxBatchSize):
+    builder = trt.Builder(TRT_LOGGER)
+    config = builder.create_builder_config()
+    engine = create_engine(maxBatchSize, builder, config, trt.float32)
+    assert engine
+    with open(ENGINE_PATH, "wb") as f:
+        f.write(engine.serialize())
+
+    del engine
+    del builder
+
+
+def doInference(context, host_in, host_out, batchSize):
+    engine = context.engine
+    assert engine.num_bindings == 2
+
+    devide_in = cuda.mem_alloc(host_in.nbytes)
+    devide_out = cuda.mem_alloc(host_out.nbytes)
+    bindings = [int(devide_in), int(devide_out)]
+    stream = cuda.Stream()
+
+    cuda.memcpy_htod_async(devide_in, host_in, stream)
+    context.execute_async(bindings=bindings, stream_handle=stream.handle)
+    cuda.memcpy_dtoh_async(host_out, devide_out, stream)
+    stream.synchronize()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-s", action='store_true')
+    parser.add_argument("-d", action='store_true')
+    args = parser.parse_args()
+
+    if not (args.s ^ args.d):
+        print(
+            "arguments not right!\n"
+            "python wide_resnet50.py -s   # serialize model to plan file\n"
+            "python wide_resnet50.py -d   # deserialize plan file and run inference"
+        )
+        sys.exit()
+
+    if args.s:
+        APIToModel(BATCH_SIZE)
+    else:
+        runtime = trt.Runtime(TRT_LOGGER)
+        assert runtime
+
+        with open(ENGINE_PATH, "rb") as f:
+            engine = runtime.deserialize_cuda_engine(f.read())
+        assert engine
+
+        context = engine.create_execution_context()
+        assert context
+
+        data = np.ones((BATCH_SIZE * 3 * INPUT_H * INPUT_W), dtype=np.float32)
+        host_in = cuda.pagelocked_empty(BATCH_SIZE * 3 * INPUT_H * INPUT_W,
+                                        dtype=np.float32)
+        np.copyto(host_in, data.ravel())
+        host_out = cuda.pagelocked_empty(OUTPUT_SIZE, dtype=np.float32)
+
+        doInference(context, host_in, host_out, BATCH_SIZE)
+
+        print(f'Output: \n{host_out[:10]}\n{host_out[-10:]}')
--- a/resnet/wideresnet50.cpp
+++ b/resnet/wideresnet50.cpp
@ -0,0 +1,367 @@
+#include "NvInfer.h"
+#include "cuda_runtime_api.h"
+#include "logging.h"
+#include <fstream>
+#include <iostream>
+#include <map>
+#include <sstream>
+#include <vector>
+#include <chrono>
+#include <cmath>
+
+#define CHECK(status) \
+    do\
+    {\
+        auto ret = (status);\
+        if (ret != 0)\
+        {\
+            std::cerr << "Cuda failure: " << ret << std::endl;\
+            abort();\
+        }\
+    } while (0)
+
+// stuff we know about the network and the input/output blobs
+static const int INPUT_H = 224;
+static const int INPUT_W = 224;
+static const int OUTPUT_SIZE = 1000;
+
+const char* INPUT_BLOB_NAME = "data";
+const char* OUTPUT_BLOB_NAME = "prob";
+
+using namespace nvinfer1;
+
+static Logger gLogger;
+
+// Load weights from files shared with TensorRT samples.
+// TensorRT weight files have a simple space delimited format:
+// [type] [size] <data x size in hex>
+std::map<std::string, Weights> loadWeights(const std::string file)
+{
+    std::cout << "Loading weights: " << file << std::endl;
+    std::map<std::string, Weights> weightMap;
+
+    // Open weights file
+    std::ifstream input(file);
+    assert(input.is_open() && "Unable to load weight file.");
+
+    // Read number of weight blobs
+    int32_t count;
+    input >> count;
+    assert(count > 0 && "Invalid weight map file.");
+
+    while (count--)
+    {
+        Weights wt{DataType::kFLOAT, nullptr, 0};
+        uint32_t size;
+
+        // Read name and type of blob
+        std::string name;
+        input >> name >> std::dec >> size;
+        wt.type = DataType::kFLOAT;
+
+        // Load blob
+        uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
+        for (uint32_t x = 0, y = size; x < y; ++x)
+        {
+            input >> std::hex >> val[x];
+        }
+        wt.values = val;
+
+        wt.count = size;
+        weightMap[name] = wt;
+    }
+
+    return weightMap;
+}
+
+IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) {
+    float *gamma = (float*)weightMap[lname + ".weight"].values;
+    float *beta = (float*)weightMap[lname + ".bias"].values;
+    float *mean = (float*)weightMap[lname + ".running_mean"].values;
+    float *var = (float*)weightMap[lname + ".running_var"].values;
+    int len = weightMap[lname + ".running_var"].count;
+    std::cout << "len " << len << std::endl;
+
+    float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
+    for (int i = 0; i < len; i++) {
+        scval[i] = gamma[i] / sqrt(var[i] + eps);
+    }
+    Weights scale{DataType::kFLOAT, scval, len};
+
+    float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
+    for (int i = 0; i < len; i++) {
+        shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
+    }
+    Weights shift{DataType::kFLOAT, shval, len};
+
+    float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
+    for (int i = 0; i < len; i++) {
+        pval[i] = 1.0;
+    }
+    Weights power{DataType::kFLOAT, pval, len};
+
+    weightMap[lname + ".scale"] = scale;
+    weightMap[lname + ".shift"] = shift;
+    weightMap[lname + ".power"] = power;
+    IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
+    assert(scale_1);
+    return scale_1;
+}
+
+IActivationLayer* bottleneck(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int inch, int outch, int stride, std::string lname) {
+    Weights emptywts{DataType::kFLOAT, nullptr, 0};
+
+    IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{1, 1}, weightMap[lname + "conv1.weight"], emptywts);
+    assert(conv1);
+
+    IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + "bn1", 1e-5);
+
+    IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
+    assert(relu1);
+
+    IConvolutionLayer* conv2 = network->addConvolutionNd(*relu1->getOutput(0), outch, DimsHW{3, 3}, weightMap[lname + "conv2.weight"], emptywts);
+    assert(conv2);
+    conv2->setStrideNd(DimsHW{stride, stride});
+    conv2->setPaddingNd(DimsHW{1, 1});
+
+    IScaleLayer* bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + "bn2", 1e-5);
+
+    IActivationLayer* relu2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU);
+    assert(relu2);
+
+    IConvolutionLayer* conv3 = network->addConvolutionNd(*relu2->getOutput(0), outch * 2, DimsHW{1, 1}, weightMap[lname + "conv3.weight"], emptywts);
+    assert(conv3);
+
+    IScaleLayer* bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + "bn3", 1e-5);
+
+    IElementWiseLayer* ew1;
+    if (stride != 1 || inch != outch * 2) {
+        IConvolutionLayer* conv4 = network->addConvolutionNd(input, outch * 2, DimsHW{1, 1}, weightMap[lname + "downsample.0.weight"], emptywts);
+        assert(conv4);
+        conv4->setStrideNd(DimsHW{stride, stride});
+
+        IScaleLayer* bn4 = addBatchNorm2d(network, weightMap, *conv4->getOutput(0), lname + "downsample.1", 1e-5);
+        ew1 = network->addElementWise(*bn4->getOutput(0), *bn3->getOutput(0), ElementWiseOperation::kSUM);
+    } else {
+        ew1 = network->addElementWise(input, *bn3->getOutput(0), ElementWiseOperation::kSUM);
+    }
+    IActivationLayer* relu3 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU);
+    assert(relu3);
+    return relu3;
+}
+
+// Create the engine using only the API and not any parser.
+ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) {
+    INetworkDefinition* network = builder->createNetworkV2(0U);
+
+    // Create input tensor of shape { 3, INPUT_H, INPUT_W } with name INPUT_BLOB_NAME
+    ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W});
+    assert(data);
+
+    std::map<std::string, Weights> weightMap = loadWeights("../wideresnet50.wts");
+    Weights emptywts{DataType::kFLOAT, nullptr, 0};
+
+    IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 64, DimsHW{7, 7}, weightMap["conv1.weight"], emptywts);
+    assert(conv1);
+    conv1->setStrideNd(DimsHW{2, 2});
+    conv1->setPaddingNd(DimsHW{3, 3});
+
+    IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "bn1", 1e-5);
+
+    // Add activation layer using the ReLU algorithm.
+    IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
+    assert(relu1);
+
+    // Add max pooling layer with stride of 2x2 and kernel size of 2x2.
+    IPoolingLayer* pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{3, 3});
+    assert(pool1);
+    pool1->setStrideNd(DimsHW{2, 2});
+    pool1->setPaddingNd(DimsHW{1, 1});
+
+    IActivationLayer* x = bottleneck(network, weightMap, *pool1->getOutput(0), 64, 128, 1, "layer1.0.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 256, 128, 1, "layer1.1.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 256, 128, 1, "layer1.2.");
+
+    x = bottleneck(network, weightMap, *x->getOutput(0), 256, 256, 2, "layer2.0.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 512, 256, 1, "layer2.1.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 512, 256, 1, "layer2.2.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 512, 256, 1, "layer2.3.");
+
+    x = bottleneck(network, weightMap, *x->getOutput(0), 512, 512, 2, "layer3.0.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 1024, 512, 1, "layer3.1.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 1024, 512, 1, "layer3.2.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 1024, 512, 1, "layer3.3.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 1024, 512, 1, "layer3.4.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 1024, 512, 1, "layer3.5.");
+
+    x = bottleneck(network, weightMap, *x->getOutput(0), 1024, 1024, 2, "layer4.0.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 2048, 1024, 1, "layer4.1.");
+    x = bottleneck(network, weightMap, *x->getOutput(0), 2048, 1024, 1, "layer4.2.");
+
+    IPoolingLayer* pool2 = network->addPoolingNd(*x->getOutput(0), PoolingType::kAVERAGE, DimsHW{7, 7});
+    assert(pool2);
+    pool2->setStrideNd(DimsHW{1, 1});
+
+    IFullyConnectedLayer* fc1 = network->addFullyConnected(*pool2->getOutput(0), 1000, weightMap["fc.weight"], weightMap["fc.bias"]);
+    assert(fc1);
+
+    fc1->getOutput(0)->setName(OUTPUT_BLOB_NAME);
+    std::cout << "set name out" << std::endl;
+    network->markOutput(*fc1->getOutput(0));
+
+    // Build engine
+    builder->setMaxBatchSize(maxBatchSize);
+    config->setMaxWorkspaceSize(1 << 20);
+    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
+    std::cout << "build out" << std::endl;
+
+    // Don't need the network any more
+    network->destroy();
+
+    // Release host memory
+    for (auto& mem : weightMap)
+    {
+        free((void*) (mem.second.values));
+    }
+
+    return engine;
+}
+
+void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream)
+{
+    // Create builder
+    IBuilder* builder = createInferBuilder(gLogger);
+    IBuilderConfig* config = builder->createBuilderConfig();
+
+    // Create model to populate the network, then set the outputs and create an engine
+    ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
+    assert(engine != nullptr);
+
+    // Serialize the engine
+    (*modelStream) = engine->serialize();
+
+    // Close everything down
+    engine->destroy();
+    builder->destroy();
+    config->destroy();
+}
+
+void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
+{
+    const ICudaEngine& engine = context.getEngine();
+
+    // Pointers to input and output device buffers to pass to engine.
+    // Engine requires exactly IEngine::getNbBindings() number of buffers.
+    assert(engine.getNbBindings() == 2);
+    void* buffers[2];
+
+    // In order to bind the buffers, we need to know the names of the input and output tensors.
+    // Note that indices are guaranteed to be less than IEngine::getNbBindings()
+    const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
+    const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
+
+    // Create GPU buffers on device
+    CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
+    CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
+
+    // Create stream
+    cudaStream_t stream;
+    CHECK(cudaStreamCreate(&stream));
+
+    // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
+    CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
+    context.enqueue(batchSize, buffers, stream, nullptr);
+    CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
+    cudaStreamSynchronize(stream);
+
+    // Release stream and buffers
+    cudaStreamDestroy(stream);
+    CHECK(cudaFree(buffers[inputIndex]));
+    CHECK(cudaFree(buffers[outputIndex]));
+}
+
+int main(int argc, char** argv)
+{
+    if (argc != 2) {
+        std::cerr << "arguments not right!" << std::endl;
+        std::cerr << "./wideresnet -s   // serialize model to plan file" << std::endl;
+        std::cerr << "./wideresnet -d   // deserialize plan file and run inference" << std::endl;
+        return -1;
+    }
+
+    // create a model using the API directly and serialize it to a stream
+    char *trtModelStream{nullptr};
+    size_t size{0};
+
+    if (std::string(argv[1]) == "-s") {
+        IHostMemory* modelStream{nullptr};
+        APIToModel(1, &modelStream);
+        assert(modelStream != nullptr);
+
+        std::ofstream p("wideresnet50.engine", std::ios::binary);
+        if (!p)
+        {
+            std::cerr << "could not open plan output file" << std::endl;
+            return -1;
+        }
+        p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
+        modelStream->destroy();
+        return 1;
+    } else if (std::string(argv[1]) == "-d") {
+        std::ifstream file("wideresnet50.engine", std::ios::binary);
+        if (file.good()) {
+            file.seekg(0, file.end);
+            size = file.tellg();
+            file.seekg(0, file.beg);
+            trtModelStream = new char[size];
+            assert(trtModelStream);
+            file.read(trtModelStream, size);
+            file.close();
+        }
+    } else {
+        return -1;
+    }
+
+
+    // Subtract mean from image
+    static float data[3 * INPUT_H * INPUT_W];
+    for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
+        data[i] = 1.0;
+
+    IRuntime* runtime = createInferRuntime(gLogger);
+    assert(runtime != nullptr);
+    ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
+    assert(engine != nullptr);
+    IExecutionContext* context = engine->createExecutionContext();
+    assert(context != nullptr);
+    delete[] trtModelStream;
+
+    // Run inference
+    static float prob[OUTPUT_SIZE];
+    for (int i = 0; i < 100; i++) {
+        auto start = std::chrono::system_clock::now();
+        doInference(*context, data, prob, 1);
+        auto end = std::chrono::system_clock::now();
+        std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
+    }
+
+    // Destroy the engine
+    context->destroy();
+    engine->destroy();
+    runtime->destroy();
+
+    // Print histogram of the output distribution
+    std::cout << "\nOutput:\n\n";
+    for (unsigned int i = 0; i < 10; i++)
+    {
+        std::cout << prob[i] << ", ";
+    }
+    std::cout << std::endl;
+    for (unsigned int i = 0; i < 10; i++)
+    {
+        std::cout << prob[OUTPUT_SIZE - 10 + i] << ", ";
+    }
+    std::cout << std::endl;
+
+    return 0;
+}