Torchvision Preprocessing Converter

The Torchvision-to-OpenVINO converter enables automatic translation of operators from the torchvision preprocessing pipeline to the OpenVINO format and embed them in your model. It is often used to adjust images serving as input for AI models to have proper dimensions or data types.

As the converter is fully based on the openvino.preprocess module, you can implement the torchvision.transforms feature easily and without the use of external libraries, reducing the overall application complexity and enabling additional performance optimizations.

Note

Not all torchvision transforms are supported yet. The following operations are available:

transforms.Compose
transforms.Normalize
transforms.ConvertImageDtype
transforms.Grayscale
transforms.Pad
transforms.ToTensor
transforms.CenterCrop
transforms.Resize

Note

transforms.Resize with InterpolationMode.NEAREST is implemented on top of the Interpolate operator’s nearest mode, configured with coordinate_transformation_mode and nearest_mode attributes chosen to closely mirror Pillow’s rounding behavior. This configuration is not guaranteed to be bit-exact with torchvision/Pillow for every input/output size combination, especially for non-square images where the height and width scale factors differ. See the nearest_mode note in the Interpolate operation specification for details.

Example

    import torch.nn.functional as f
    import openvino as ov
    import numpy as np
    import torchvision
    import torch
    import os

    from openvino.preprocess.torchvision import PreprocessConverter
    from PIL import Image


    # 1. Create a sample model
    class Convnet(torch.nn.Module):
        def __init__(self, input_channels):
            super(Convnet, self).__init__()
            self.conv1 = torch.nn.Conv2d(input_channels, 6, 5)
            self.conv2 = torch.nn.Conv2d(6, 16, 3)

        def forward(self, data):
            data = f.max_pool2d(f.relu(self.conv1(data)), 2)
            data = f.max_pool2d(f.relu(self.conv2(data)), 2)
            return data


    # 2. Define torchvision preprocessing pipeline
    preprocess_pipeline = torchvision.transforms.Compose(
       [
           torchvision.transforms.Resize(256, interpolation=torchvision.transforms.InterpolationMode.NEAREST),
           torchvision.transforms.CenterCrop((216, 218)),
           torchvision.transforms.Pad((2, 3, 4, 5), fill=3),
           torchvision.transforms.ToTensor(),
           torchvision.transforms.ConvertImageDtype(torch.float32),
           torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
       ]
   )

    # 3. Read the model into OpenVINO
    torch_model = Convnet(input_channels=3)
    torch.onnx.export(torch_model, torch.randn(1, 3, 224, 224), "test_convnet.onnx", verbose=False, input_names=["input"], output_names=["output"])
    core = ov.Core()
    ov_model = core.read_model(model="test_convnet.onnx")
    if os.path.exists("test_convnet.onnx"):
        os.remove("test_convnet.onnx")
    test_input = np.random.randint(255, size=(260, 260, 3), dtype=np.uint16)

    # 4. Embed the torchvision preocessing into OpenVINO model
    ov_model = PreprocessConverter.from_torchvision(
        model=ov_model, transform=preprocess_pipeline, input_example=Image.fromarray(test_input.astype("uint8"), "RGB")
    )
    ov_model = core.compile_model(ov_model, "CPU")

    # 5. Perform inference
    ov_input = np.expand_dims(test_input, axis=0)
    output = ov_model.output(0)
    ov_result = ov_model(ov_input)[output]