chore: import upstream snapshot with attribution

samples/python/network_api_pytorch_mnist/README.md

@@ -0,0 +1,118 @@
+# “Hello World” For TensorRT Using PyTorch And Python
+
+**Table Of Contents**
+
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+    * [TensorRT API layers and ops](#tensorrt-api-layers-and-ops)
+- [Prerequisites](#prerequisites)
+- [Running the sample](#running-the-sample)
+    * [Sample `--help` options](#sample-help-options)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample, `network_api_pytorch_mnist`, trains a convolutional model on the [MNIST](https://ossci-datasets.s3.amazonaws.com/mnist/) dataset and runs inference with a TensorRT engine.
+
+## How does this sample work?
+
+This sample is an end-to-end sample that trains a model in PyTorch, recreates the network in TensorRT, imports weights from the trained model, and finally runs inference with a TensorRT engine. For more information, see [Creating A Network Definition In Python](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#network_python).
+
+The `sample.py` script imports the functions from the `mnist.py` script for training the PyTorch model, as well as retrieving test cases from the PyTorch Data Loader.
+
+### TensorRT API layers and ops
+
+In this sample, the following layers are used. For more information about these layers, see the [TensorRT Developer Guide: Layers](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#layers) documentation.
+
+[Activation layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#activation-layer)
+The Activation layer implements element-wise activation functions. Specifically, this sample uses the Activation layer with the type `RELU`.
+
+[Convolution layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#convolution-layer)
+The Convolution layer computes a 2D (channel, height, and width) convolution, with or without bias.
+
+[MatrixMultiplyLayer](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#matrixmultiply-layer)
+The MatrixMultiply layer implements a matrix multiplication.
+(The [FullyConnected layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#fullyconnected-layer) is deprecated since 8.4.
+The bias of FullyConnected semantic can be added with an
+[ElementwiseLayer](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#elementwise-layer) of `SUM` operation.)
+
+[Pooling layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#pooling-layer)
+The Pooling layer implements pooling within a channel. Supported pooling types are `maximum`, `average` and `maximum-average blend`.
+
+## Prerequisites
+
+1. Upgrade pip version and install the sample dependencies.
+    ```bash
+    pip3 install --upgrade pip
+    pip3 install -r requirements.txt
+    ```
+
+To run this sample you must be using Python 3.6 or newer.
+
+On PowerPC systems, you will need to manually install PyTorch using IBM's [PowerAI](https://www.ibm.com/support/knowledgecenter/SS5SF7_1.6.0/navigation/pai_install.htm).
+
+2. Preparing sample data
+
+See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README.
+
+The MNIST dataset can be found under `$TRT_DATADIR/mnist`.
+
+## Running the sample
+
+1.  Run the sample to create a TensorRT inference engine and run inference:
+    `python3 sample.py`
+
+2.  Verify that the sample ran successfully. If the sample runs successfully you should see a match between the test case and the prediction.
+     ```
+    Test Case: 0
+    Prediction: 0
+    ```
+
+### Sample --help options
+
+To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
+
+# Additional resources
+
+The following resources provide a deeper understanding about getting started with TensorRT using Python:
+
+**Model**
+- [MNIST model](https://github.com/pytorch/examples/tree/master/mnist)
+
+**Dataset**
+- [MNIST database](https://ossci-datasets.s3.amazonaws.com/mnist/)
+
+**Documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.
+
+August 2023
+Removed support for Python versions < 3.8.
+
+September 2021
+Updated the sample to use explicit batch network definition.
+
+March 2021
+Documented the Python version limitations.
+
+February 2019
+This `README.md` file was recreated, updated and reviewed.
+
+# Known issues
+
+This sample only supports Python 3.6+ due to `torch` and `torchvision` version requirements.

samples/python/network_api_pytorch_mnist/model.py

@@ -0,0 +1,156 @@
+#
+# 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.
+#
+
+# This file contains functions for training a PyTorch MNIST Model
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.autograd import Variable
+
+import numpy as np
+
+from random import randint
+
+
+# Network
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 20, kernel_size=5)
+        self.conv2 = nn.Conv2d(20, 50, kernel_size=5)
+        self.fc1 = nn.Linear(800, 500)
+        self.fc2 = nn.Linear(500, 10)
+
+    def forward(self, x):
+        x = F.max_pool2d(self.conv1(x), kernel_size=2, stride=2)
+        x = F.max_pool2d(self.conv2(x), kernel_size=2, stride=2)
+        x = x.view(-1, 800)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+class MnistModel(object):
+    def __init__(self):
+        self.batch_size = 64
+        self.test_batch_size = 100
+        self.learning_rate = 0.0025
+        self.sgd_momentum = 0.9
+        self.log_interval = 100
+        # Fetch MNIST data set.
+        self.train_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "/tmp/mnist/data",
+                train=True,
+                download=True,
+                transform=transforms.Compose(
+                    [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+                ),
+            ),
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=1,
+            timeout=600,
+        )
+        self.test_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "/tmp/mnist/data",
+                train=False,
+                transform=transforms.Compose(
+                    [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+                ),
+            ),
+            batch_size=self.test_batch_size,
+            shuffle=True,
+            num_workers=1,
+            timeout=600,
+        )
+        self.network = Net()
+        if torch.cuda.is_available():
+            self.network = self.network.to("cuda")
+
+    # Train the network for one or more epochs, validating after each epoch.
+    def learn(self, num_epochs=2):
+        # Train the network for a single epoch
+        def train(epoch):
+            self.network.train()
+            optimizer = optim.SGD(
+                self.network.parameters(),
+                lr=self.learning_rate,
+                momentum=self.sgd_momentum,
+            )
+            for batch, (data, target) in enumerate(self.train_loader):
+                if torch.cuda.is_available():
+                    data = data.to("cuda")
+                    target = target.to("cuda")
+                data, target = Variable(data), Variable(target)
+                optimizer.zero_grad()
+                output = self.network(data)
+                loss = F.nll_loss(output, target)
+                loss.backward()
+                optimizer.step()
+                if batch % self.log_interval == 0:
+                    print(
+                        "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                            epoch,
+                            batch * len(data),
+                            len(self.train_loader.dataset),
+                            100.0 * batch / len(self.train_loader),
+                            loss.data.item(),
+                        )
+                    )
+
+        # Test the network
+        def test(epoch):
+            self.network.eval()
+            test_loss = 0
+            correct = 0
+            for data, target in self.test_loader:
+                with torch.no_grad():
+                    if torch.cuda.is_available():
+                        data = data.to("cuda")
+                        target = target.to("cuda")
+                    data, target = Variable(data), Variable(target)
+                output = self.network(data)
+                test_loss += F.nll_loss(output, target).data.item()
+                pred = output.data.max(1)[1]
+                correct += pred.eq(target.data).cpu().sum()
+            test_loss /= len(self.test_loader)
+            print(
+                "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(
+                    test_loss,
+                    correct,
+                    len(self.test_loader.dataset),
+                    100.0 * correct / len(self.test_loader.dataset),
+                )
+            )
+
+        for e in range(num_epochs):
+            train(e + 1)
+            test(e + 1)
+
+    def get_weights(self):
+        return self.network.state_dict()
+
+    def get_random_testcase(self):
+        data, target = next(iter(self.test_loader))
+        case_num = randint(0, len(data) - 1)
+        test_case = data.cpu().numpy()[case_num].ravel().astype(np.float32)
+        test_name = target.cpu().numpy()[case_num]
+        return test_case, test_name

samples/python/network_api_pytorch_mnist/requirements.txt

@@ -0,0 +1,9 @@
+Pillow==11.3.0
+torch
+torchvision
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/network_api_pytorch_mnist/sample.py

@@ -0,0 +1,179 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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.
+#
+
+import os
+import sys
+
+# This sample uses an MNIST PyTorch model to create a TensorRT Inference Engine
+import model
+import numpy as np
+
+import tensorrt as trt
+
+sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
+import common
+
+# You can set the logger severity higher to suppress messages (or lower to display more messages).
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+
+
+class ModelData(object):
+    INPUT_NAME = "data"
+    INPUT_SHAPE = (1, 1, 28, 28)
+    OUTPUT_NAME = "prob"
+    OUTPUT_SIZE = 10
+    DTYPE = trt.float32
+
+
+def populate_network(network, weights):
+    # Configure the network layers based on the weights provided.
+    input_tensor = network.add_input(
+        name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE
+    )
+
+    def add_matmul_as_fc(net, input, outputs, w, b):
+        assert len(input.shape) >= 3
+        m = 1 if len(input.shape) == 3 else input.shape[0]
+        k = int(np.prod(input.shape) / m)
+        assert np.prod(input.shape) == m * k
+        n = int(w.size / k)
+        assert w.size == n * k
+        assert b.size == n
+
+        input_reshape = net.add_shuffle(input)
+        input_reshape.reshape_dims = trt.Dims2(m, k)
+
+        filter_const = net.add_constant(trt.Dims2(n, k), w)
+        mm = net.add_matrix_multiply(
+            input_reshape.get_output(0),
+            trt.MatrixOperation.NONE,
+            filter_const.get_output(0),
+            trt.MatrixOperation.TRANSPOSE,
+        )
+
+        bias_const = net.add_constant(trt.Dims2(1, n), b)
+        bias_add = net.add_elementwise(
+            mm.get_output(0), bias_const.get_output(0), trt.ElementWiseOperation.SUM
+        )
+
+        output_reshape = net.add_shuffle(bias_add.get_output(0))
+        output_reshape.reshape_dims = trt.Dims4(m, n, 1, 1)
+        return output_reshape
+
+    conv1_w = weights["conv1.weight"].cpu().numpy()
+    conv1_b = weights["conv1.bias"].cpu().numpy()
+    conv1 = network.add_convolution_nd(
+        input=input_tensor,
+        num_output_maps=20,
+        kernel_shape=(5, 5),
+        kernel=conv1_w,
+        bias=conv1_b,
+    )
+    conv1.stride_nd = (1, 1)
+
+    pool1 = network.add_pooling_nd(
+        input=conv1.get_output(0), type=trt.PoolingType.MAX, window_size=(2, 2)
+    )
+    pool1.stride_nd = trt.Dims2(2, 2)
+
+    conv2_w = weights["conv2.weight"].cpu().numpy()
+    conv2_b = weights["conv2.bias"].cpu().numpy()
+    conv2 = network.add_convolution_nd(
+        pool1.get_output(0), 50, (5, 5), conv2_w, conv2_b
+    )
+    conv2.stride_nd = (1, 1)
+
+    pool2 = network.add_pooling_nd(conv2.get_output(0), trt.PoolingType.MAX, (2, 2))
+    pool2.stride_nd = trt.Dims2(2, 2)
+
+    fc1_w = weights["fc1.weight"].cpu().numpy()
+    fc1_b = weights["fc1.bias"].cpu().numpy()
+    fc1 = add_matmul_as_fc(network, pool2.get_output(0), 500, fc1_w, fc1_b)
+
+    relu1 = network.add_activation(
+        input=fc1.get_output(0), type=trt.ActivationType.RELU
+    )
+
+    fc2_w = weights["fc2.weight"].cpu().numpy()
+    fc2_b = weights["fc2.bias"].cpu().numpy()
+    fc2 = add_matmul_as_fc(
+        network, relu1.get_output(0), ModelData.OUTPUT_SIZE, fc2_w, fc2_b
+    )
+
+    fc2.get_output(0).name = ModelData.OUTPUT_NAME
+    network.mark_output(tensor=fc2.get_output(0))
+
+
+def build_engine(weights):
+    # For more information on TRT basics, refer to the introductory samples.
+    builder = trt.Builder(TRT_LOGGER)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    config = builder.create_builder_config()
+    runtime = trt.Runtime(TRT_LOGGER)
+
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, common.GiB(1))
+    # Populate the network using weights from the PyTorch model.
+    populate_network(network, weights)
+    # Build and return an engine.
+    plan = builder.build_serialized_network(network, config)
+    return runtime.deserialize_cuda_engine(plan)
+
+
+# Loads a random test case from pytorch's DataLoader
+def load_random_test_case(model, pagelocked_buffer):
+    # Select an image at random to be the test case.
+    img, expected_output = model.get_random_testcase()
+    # Copy to the pagelocked input buffer
+    np.copyto(pagelocked_buffer, img)
+    return expected_output
+
+
+def main():
+    common.add_help(description="Runs an MNIST network using a PyTorch model")
+    # Train the PyTorch model
+    mnist_model = model.MnistModel()
+    mnist_model.learn()
+    weights = mnist_model.get_weights()
+    # Do inference with TensorRT.
+    engine = build_engine(weights)
+
+    # Build an engine, allocate buffers and create a stream.
+    # For more information on buffer allocation, refer to the introductory samples.
+    inputs, outputs, bindings = common.allocate_buffers(engine)
+    context = engine.create_execution_context()
+
+    # Use context manager for proper stream lifecycle management
+    with common.CudaStreamContext() as stream:
+        case_num = load_random_test_case(mnist_model, pagelocked_buffer=inputs[0].host)
+        # For more information on performing inference, refer to the introductory samples.
+        # The common.do_inference function will return a list of outputs - we only have one in this case.
+        [output] = common.do_inference(
+            context,
+            engine=engine,
+            bindings=bindings,
+            inputs=inputs,
+            outputs=outputs,
+            stream=stream,
+        )
+        pred = np.argmax(output)
+        common.free_buffers(inputs, outputs)
+    print("Test Case: " + str(case_num))
+    print("Prediction: " + str(pred))
+
+
+if __name__ == "__main__":
+    main()