chore: import upstream snapshot with attribution

quickstart/SemanticSegmentation/.gitignore

@@ -0,0 +1,5 @@
+bin/*
+.ipynb_checkpoints/
+*.ppm
+*.onnx
+*.engine

quickstart/SemanticSegmentation/Makefile

@@ -0,0 +1,22 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+OUTNAME_RELEASE = segmentation_tutorial
+OUTNAME_DEBUG   = segmentation_tutorial_debug
+EXTRA_DIRECTORIES = ../common
+MAKEFILE ?= ../Makefile.config
+include $(MAKEFILE)

quickstart/SemanticSegmentation/export.py

@@ -0,0 +1,70 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from PIL import Image
+from io import BytesIO
+import requests
+
+output_image = "input.ppm"
+
+# Read sample image input and save it in ppm format
+print("Exporting ppm image {}".format(output_image))
+response = requests.get("https://pytorch.org/assets/images/deeplab1.png")
+with Image.open(BytesIO(response.content)) as img:
+    ppm = Image.new("RGB", img.size, (255, 255, 255))
+    ppm.paste(img, mask=img.split()[3])
+    ppm.save(output_image)
+
+
+import torch
+import torch.nn as nn
+import torchvision.models.segmentation as segmentation
+
+output_onnx = "fcn-resnet101.onnx"
+
+
+# FC-ResNet101 pretrained model from torch-hub extended with argmax layer
+class FCN_ResNet101(nn.Module):
+    def __init__(self):
+        super(FCN_ResNet101, self).__init__()
+        self.model = segmentation.fcn_resnet101(pretrained=True)
+
+    def forward(self, inputs):
+        x = self.model(inputs)["out"]
+        x = x.argmax(1, keepdims=True)
+        return x
+
+
+model = FCN_ResNet101()
+model.eval()
+
+# Generate input tensor with random values
+input_tensor = torch.rand(4, 3, 224, 224)
+
+# Export torch model to ONNX
+print("Exporting ONNX model {}".format(output_onnx))
+torch.onnx.export(
+    model,
+    input_tensor,
+    output_onnx,
+    opset_version=12,
+    do_constant_folding=True,
+    input_names=["input"],
+    output_names=["output"],
+    dynamic_axes={"input": {0: "batch", 2: "height", 3: "width"}, "output": {0: "batch", 2: "height", 3: "width"}},
+    verbose=False,
+)

quickstart/SemanticSegmentation/tutorial-runtime.cpp

@@ -0,0 +1,193 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cassert>
+#include <cfloat>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <sstream>
+
+#include <cuda_runtime_api.h>
+#include "NvInfer.h"
+#include "NvOnnxParser.h"
+#include "logger.h"
+#include "util.h"
+
+constexpr long long operator"" _MiB(long long unsigned val)
+{
+    return val * (1 << 20);
+}
+
+using sample::gLogError;
+using sample::gLogInfo;
+
+//!
+//! \class SampleSegmentation
+//!
+//! \brief Implements semantic segmentation using FCN-ResNet101 ONNX model.
+//!
+class SampleSegmentation
+{
+
+public:
+    SampleSegmentation(const std::string& engineFilename);
+    bool infer(const std::string& input_filename, int32_t width, int32_t height, const std::string& output_filename);
+
+private:
+    std::string mEngineFilename;                    //!< Filename of the serialized engine.
+
+    nvinfer1::Dims mInputDims;                      //!< The dimensions of the input to the network.
+    nvinfer1::Dims mOutputDims;                     //!< The dimensions of the output to the network.
+
+    std::unique_ptr<nvinfer1::IRuntime> mRuntime;   //!< The TensorRT runtime used to run the network
+    std::unique_ptr<nvinfer1::ICudaEngine> mEngine; //!< The TensorRT engine used to run the network
+};
+
+SampleSegmentation::SampleSegmentation(const std::string& engineFilename)
+    : mEngineFilename(engineFilename)
+    , mEngine(nullptr)
+{
+    // De-serialize engine from file
+    std::ifstream engineFile(engineFilename, std::ios::binary);
+    if (engineFile.fail())
+    {
+        return;
+    }
+
+    engineFile.seekg(0, std::ifstream::end);
+    auto fsize = engineFile.tellg();
+    engineFile.seekg(0, std::ifstream::beg);
+
+    std::vector<char> engineData(fsize);
+    engineFile.read(engineData.data(), fsize);
+
+    mRuntime.reset(nvinfer1::createInferRuntime(sample::gLogger.getTRTLogger()));
+    mEngine.reset(mRuntime->deserializeCudaEngine(engineData.data(), fsize));
+    assert(mEngine.get() != nullptr);
+}
+
+//!
+//! \brief Runs the TensorRT inference.
+//!
+//! \details Allocate input and output memory, and executes the engine.
+//!
+bool SampleSegmentation::infer(const std::string& input_filename, int32_t width, int32_t height, const std::string& output_filename)
+{
+    auto context = std::unique_ptr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
+    if (!context)
+    {
+        return false;
+    }
+
+    char const* input_name = "input";
+    assert(mEngine->getTensorDataType(input_name) == nvinfer1::DataType::kFLOAT);
+    auto input_dims = nvinfer1::Dims4{1, /* channels */ 3, height, width};
+    context->setInputShape(input_name, input_dims);
+    auto input_size = util::getMemorySize(input_dims, sizeof(float));
+
+    char const* output_name = "output";
+    assert(mEngine->getTensorDataType(output_name) == nvinfer1::DataType::kINT64);
+    auto output_dims = context->getTensorShape(output_name);
+    auto output_size = util::getMemorySize(output_dims, sizeof(int64_t));
+
+    // Allocate CUDA memory for input and output bindings
+    void* input_mem{nullptr};
+    if (cudaMalloc(&input_mem, input_size) != cudaSuccess)
+    {
+        gLogError << "ERROR: input cuda memory allocation failed, size = " << input_size << " bytes" << std::endl;
+        return false;
+    }
+    void* output_mem{nullptr};
+    if (cudaMalloc(&output_mem, output_size) != cudaSuccess)
+    {
+        gLogError << "ERROR: output cuda memory allocation failed, size = " << output_size << " bytes" << std::endl;
+        return false;
+    }
+
+    // Read image data from file and mean-normalize it
+    const std::vector<float> mean{0.485f, 0.456f, 0.406f};
+    const std::vector<float> stddev{0.229f, 0.224f, 0.225f};
+    auto input_image{util::RGBImageReader(input_filename, input_dims, mean, stddev)};
+    input_image.read();
+    auto input_buffer = input_image.process();
+    cudaStream_t stream;
+    if (cudaStreamCreate(&stream) != cudaSuccess)
+    {
+        gLogError << "ERROR: cuda stream creation failed." << std::endl;
+        return false;
+    }
+
+    // Copy image data to input binding memory
+    if (cudaMemcpyAsync(input_mem, input_buffer.get(), input_size, cudaMemcpyHostToDevice, stream) != cudaSuccess)
+    {
+        gLogError << "ERROR: CUDA memory copy of input failed, size = " << input_size << " bytes" << std::endl;
+        return false;
+    }
+    context->setTensorAddress(input_name, input_mem);
+    context->setTensorAddress(output_name, output_mem);
+
+    // Run TensorRT inference
+    bool status = context->enqueueV3(stream);
+    if (!status)
+    {
+        gLogError << "ERROR: TensorRT inference failed" << std::endl;
+        return false;
+    }
+
+    // Copy predictions from output binding memory
+    auto output_buffer = std::unique_ptr<int64_t>{new int64_t[output_size]};
+    if (cudaMemcpyAsync(output_buffer.get(), output_mem, output_size, cudaMemcpyDeviceToHost, stream) != cudaSuccess)
+    {
+        gLogError << "ERROR: CUDA memory copy of output failed, size = " << output_size << " bytes" << std::endl;
+        return false;
+    }
+    cudaStreamSynchronize(stream);
+
+    // Plot the semantic segmentation predictions of 21 classes in a colormap image and write to file
+    const int num_classes{21};
+    const std::vector<int> palette{(0x1 << 25) - 1, (0x1 << 15) - 1, (0x1 << 21) - 1};
+    auto output_image{util::ArgmaxImageWriter(output_filename, output_dims, palette, num_classes)};
+    int64_t* output_ptr = output_buffer.get();
+    std::vector<int32_t> output_buffer_casted(output_size);
+    for (size_t i = 0; i < output_size; ++i) {
+        output_buffer_casted[i] = static_cast<int32_t>(output_ptr[i]);
+    }
+    output_image.process(output_buffer_casted.data());
+    output_image.write();
+
+    // Free CUDA resources
+    cudaFree(input_mem);
+    cudaFree(output_mem);
+    return true;
+}
+
+int main(int argc, char** argv)
+{
+    int32_t width{1282};
+    int32_t height{1026};
+
+    SampleSegmentation sample("fcn-resnet101.engine");
+
+    gLogInfo << "Running TensorRT inference for FCN-ResNet101" << std::endl;
+    if (!sample.infer("input.ppm", width, height, "output.ppm"))
+    {
+        return -1;
+    }
+
+    return 0;
+}

quickstart/SemanticSegmentation/tutorial-runtime.ipynb

@@ -0,0 +1,341 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#\n",
+    "# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License.\n",
+    "#"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Check the TensorRT version"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!python3 -c 'import tensorrt; print(\"TensorRT version: {}\".format(tensorrt.__version__))'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Prepare the input image and ONNX model file"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!python3 /workspace/TensorRT/quickstart/SemanticSegmentation/export.py"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Build TensorRT engine from the ONNX model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!trtexec --onnx=fcn-resnet101.onnx --saveEngine=fcn-resnet101.engine --optShapes=input:1x3x1026x1282 --stronglyTyped"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Import required modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import os\n",
+    "import ctypes\n",
+    "from cuda.bindings import runtime as cudart\n",
+    "import tensorrt as trt\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "from PIL import Image\n",
+    "\n",
+    "TRT_LOGGER = trt.Logger()\n",
+    "\n",
+    "assert cudart.cudaSetDevice(0) == (cudart.cudaError_t.cudaSuccess,)\n",
+    "\n",
+    "# Filenames of TensorRT plan file and input/output images.\n",
+    "engine_file = \"/workspace/fcn-resnet101.engine\"\n",
+    "input_file  = \"/workspace/input.ppm\"\n",
+    "output_file = \"/workspace/output.ppm\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Utilities for input / output processing"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# For torchvision models, input images are loaded in to a range of [0, 1] and\n",
+    "# normalized using mean = [0.485, 0.456, 0.406] and stddev = [0.229, 0.224, 0.225].\n",
+    "def preprocess(image):\n",
+    "    # Mean normalization\n",
+    "    mean = np.array([0.485, 0.456, 0.406]).astype('float32')\n",
+    "    stddev = np.array([0.229, 0.224, 0.225]).astype('float32')\n",
+    "    data = (np.asarray(image).astype('float32') / float(255.0) - mean) / stddev\n",
+    "    # Switch from HWC to to CHW order\n",
+    "    return np.moveaxis(data, 2, 0)\n",
+    "\n",
+    "def postprocess(data):\n",
+    "    num_classes = 21\n",
+    "    # create a color palette, selecting a color for each class\n",
+    "    palette = np.array([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])\n",
+    "    colors = np.array([palette*i%255 for i in range(num_classes)]).astype(\"uint8\")\n",
+    "    # plot the segmentation predictions for 21 classes in different colors\n",
+    "    img = Image.fromarray(data.astype('uint8'), mode='P')\n",
+    "    img.putpalette(colors)\n",
+    "    return img\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Load TensorRT engine"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Deserialize the TensorRT engine from specified plan file. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_engine(engine_file_path):\n",
+    "    assert os.path.exists(engine_file_path)\n",
+    "    print(\"Reading engine from file {}\".format(engine_file_path))\n",
+    "    with open(engine_file_path, \"rb\") as f, trt.Runtime(TRT_LOGGER) as runtime:\n",
+    "        return runtime.deserialize_cuda_engine(f.read())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Inference pipeline"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Starting with a deserialized engine, TensorRT inference pipeline consists of the following steps:\n",
+    "- Create an execution context and specify input shape (based on the image dimensions for inference).\n",
+    "- Allocate CUDA device memory for input and output.\n",
+    "- Allocate CUDA page-locked host memory to efficiently copy back the output.\n",
+    "- Transfer the processed image data into input memory using asynchronous host-to-device CUDA copy.\n",
+    "- Kickoff the TensorRT inference pipeline using the asynchronous execute API.\n",
+    "- Transfer the segmentation output back into pagelocked host memory using device-to-host CUDA copy.\n",
+    "- Synchronize the stream used for data transfers and inference execution to ensure all operations are completes.\n",
+    "- Finally, write out the segmentation output to an image file for visualization."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def infer(engine, input_file, output_file):\n",
+    "    print(\"Reading input image from file {}\".format(input_file))\n",
+    "    with Image.open(input_file) as img:\n",
+    "        input_image = preprocess(img)\n",
+    "        image_width = img.width\n",
+    "        image_height = img.height\n",
+    "\n",
+    "    with engine.create_execution_context() as context:\n",
+    "        input_buffers = {}\n",
+    "        input_memories = {}\n",
+    "\n",
+    "        # Allocate host and device buffers\n",
+    "        tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]\n",
+    "        for tensor in tensor_names:\n",
+    "            size = trt.volume(context.get_tensor_shape(tensor))\n",
+    "            dtype = trt.nptype(engine.get_tensor_dtype(tensor))\n",
+    "\n",
+    "            if engine.get_tensor_mode(tensor) == trt.TensorIOMode.INPUT:\n",
+    "                context.set_input_shape(tensor, (1, 3, image_height, image_width))\n",
+    "                input_buffers[tensor] = np.ascontiguousarray(input_image)\n",
+    "                err, input_memories[tensor] = cudart.cudaMalloc(input_image.nbytes)\n",
+    "                assert err == cudart.cudaError_t.cudaSuccess\n",
+    "                context.set_tensor_address(tensor, input_memories[tensor])\n",
+    "            else:\n",
+    "                err, output_buffer_ptr = cudart.cudaMallocHost(size * dtype().itemsize)\n",
+    "                assert err == cudart.cudaError_t.cudaSuccess\n",
+    "                pointer_type = ctypes.POINTER(np.ctypeslib.as_ctypes_type(dtype))\n",
+    "                output_buffer = np.ctypeslib.as_array(ctypes.cast(output_buffer_ptr, pointer_type), (size,))\n",
+    "\n",
+    "                err, output_memory = cudart.cudaMalloc(output_buffer.nbytes)\n",
+    "                assert err == cudart.cudaError_t.cudaSuccess\n",
+    "                context.set_tensor_address(tensor, output_memory)\n",
+    "\n",
+    "        err, stream = cudart.cudaStreamCreate()\n",
+    "        assert err == cudart.cudaError_t.cudaSuccess\n",
+    "\n",
+    "        # Transfer input data to the GPU for all input tensors\n",
+    "        for tensor_name, input_buffer in input_buffers.items():\n",
+    "            input_memory = input_memories[tensor_name]\n",
+    "            err, = cudart.cudaMemcpyAsync(input_memory, input_buffer.ctypes.data, input_buffer.nbytes,\n",
+    "                                          cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream)\n",
+    "            assert err == cudart.cudaError_t.cudaSuccess\n",
+    "\n",
+    "        # Run inference\n",
+    "        context.execute_async_v3(stream)\n",
+    "\n",
+    "        # Transfer prediction output from the GPU.\n",
+    "        err, = cudart.cudaMemcpyAsync(output_buffer.ctypes.data, output_memory, output_buffer.nbytes,\n",
+    "                                      cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)\n",
+    "        assert err == cudart.cudaError_t.cudaSuccess\n",
+    "        # Synchronize the stream\n",
+    "        err, = cudart.cudaStreamSynchronize(stream)\n",
+    "        assert err == cudart.cudaError_t.cudaSuccess\n",
+    "\n",
+    "        output_d64 = np.array(output_buffer, dtype=np.int64)\n",
+    "        np.savetxt('test.out', output_d64.astype(int), fmt='%i', delimiter=' ', newline=' ')\n",
+    "\n",
+    "        with postprocess(np.reshape(output_buffer, (image_height, image_width))) as img:\n",
+    "            print(\"Writing output image to file {}\".format(output_file))\n",
+    "            img.convert('RGB').save(output_file, \"PPM\")\n",
+    "\n",
+    "        # cleanup cuda resources for all input tensors\n",
+    "        for input_memory in input_memories.values():\n",
+    "            cudart.cudaFree(input_memory)\n",
+    "        cudart.cudaFree(output_memory)\n",
+    "        cudart.cudaFreeHost(output_buffer_ptr)\n",
+    "        cudart.cudaStreamDestroy(stream)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Plot input image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.imshow(Image.open(input_file))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Run inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"Running TensorRT inference for FCN-ResNet101\")\n",
+    "with load_engine(engine_file) as engine:\n",
+    "    infer(engine, input_file, output_file)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Plot segmentation output"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.imshow(Image.open(output_file))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}