Classification Inference Using PyTorch

The classification sample in CVCUDA uses the ResNet50 deep learning model from the torchvision library. Since the model does not come with the softmax layer at the end, we are going to add one. The following code snippet shows how the model is setup for inference use case with PyTorch.

 1class ClassificationPyTorch:  # noqa: E302
 2    def __init__(
 3        self,
 4        output_dir,
 5        batch_size,
 6        image_size,
 7        device_id,
 8        cvcuda_perf,
 9    ):
10        self.logger = logging.getLogger(__name__)
11        self.output_dir = output_dir
12        self.device_id = device_id
13        self.cvcuda_perf = cvcuda_perf
14        # The underlying PyTorch model that we use for inference is the ResNet50 model
15        # from torchvision.
16        torch_model = torchvision_models.resnet50
17        weights = torchvision_models.ResNet50_Weights.DEFAULT
18        self.labels = weights.meta["categories"]
19        # Save the list of labels so that the C++ sample can read it.
20        with open(os.path.join(output_dir, "labels.txt"), "w") as f:
21            for line in self.labels:
22                f.write("%s\n" % line)
23
24        # Inference uses PyTorch to run a classification model on the pre-processed
25        # input and outputs the classification scores.
26        class Resnet50_Softmax(torch.nn.Module):
27            def __init__(self, resnet50):
28                super(Resnet50_Softmax, self).__init__()
29                self.resnet50 = resnet50
30
31            def forward(self, x):
32                infer_output = self.resnet50(x)
33                return torch.nn.functional.softmax(infer_output, dim=1)
34
35        resnet_base = torch_model(weights=weights)
36        resnet_base.eval()
37        self.model = Resnet50_Softmax(resnet_base).cuda(self.device_id)
38        self.model.eval()
39
40        self.logger.info("Using PyTorch as the inference engine.")

To run the inference the __call__ method is used. It makes sure to use the CUDA stream and perform the forward inference pass without computing gradients.

 1def __call__(self, tensor):
 2    self.cvcuda_perf.push_range("inference.torch")
 3
 4    with torch.no_grad():
 5
 6        if isinstance(tensor, torch.Tensor):
 7            if not tensor.is_cuda:
 8                tensor = tensor.to("cuda:%d" % self.device_id)
 9        else:
10            # Convert CVCUDA tensor to Torch tensor.
11            tensor = torch.as_tensor(
12                tensor.cuda(), device="cuda:%d" % self.device_id
13            )
14
15        classification_scores = self.model(tensor)
16
17    self.cvcuda_perf.pop_range()
18    return classification_scores
19