Skip to content

torch

torch

Attributes

__all__ module-attribute 

__all__ = ['ClassificationPipeline']

Classes

ClassificationPipeline 

ClassificationPipeline(runtime: Runtime[Tensor, Tensor], image_size: tuple[int, int] = (224, 224), mean: tuple[float, float, float] = (0.485, 0.456, 0.406), std: tuple[float, float, float] = (0.229, 0.224, 0.225), class_names: dict[int, str] | None = None, batch_strategy: BatchStrategy[Tensor, Tensor] | None = None)

Bases: ClassificationMixin, Pipeline[Tensor, Tensor, ClassificationOutput]

Image classification pipeline (async version).

Performs
  • Image decoding and conversion
  • Resizing and normalization
  • Model inference
  • Class prediction with confidence

Parameters:

Name Type Description Default
runtime Runtime[Tensor, Tensor]

Inference runtime.

 required
image_size tuple[int, int]

Target image size (default: (224, 224)).

 (224, 224)
mean tuple[float, float, float]

Normalization mean (default: ImageNet mean).

 (0.485, 0.456, 0.406)
std tuple[float, float, float]

Normalization std (default: ImageNet std).

 (0.229, 0.224, 0.225)
class_names dict[int, str] | None

Optional mapping from class ID to class name.

 None
batch_strategy BatchStrategy[Tensor, Tensor] | None

Optional batching strategy.

 None
Example 
runtime = TorchScriptRuntime(
    model_path="resnet50.pt", device="cuda"
)
pipeline = ClassificationPipeline(
    runtime=runtime,
    class_names={0: "cat", 1: "dog", 2: "bird"},
)

async with pipeline.serve():
    result = await pipeline(image_bytes)
    print(f"{result.class_name}: {result.confidence:.2%}")
Source code in inferflow/asyncio/pipeline/classification/torch.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def __init__(
    self,
    runtime: Runtime[torch.Tensor, torch.Tensor],
    image_size: tuple[int, int] = (224, 224),
    mean: tuple[float, float, float] = (0.485, 0.456, 0.406),
    std: tuple[float, float, float] = (0.229, 0.224, 0.225),
    class_names: dict[int, str] | None = None,
    batch_strategy: BatchStrategy[torch.Tensor, torch.Tensor] | None = None,
):
    super().__init__(runtime=runtime, batch_strategy=batch_strategy)

    self.image_size = image_size
    self.mean = torch.tensor(mean).view(3, 1, 1)
    self.std = torch.tensor(std).view(3, 1, 1)
    self.class_names = class_names or {}

    import torchvision.transforms as T

    self.transform = T.Compose([T.Resize(image_size), T.ToTensor()])
Attributes
image_size instance-attribute 
image_size = image_size
mean instance-attribute 
mean = view(3, 1, 1)
std instance-attribute 
std = view(3, 1, 1)
class_names instance-attribute 
class_names = class_names or {}
transform instance-attribute 
transform = Compose([Resize(image_size), ToTensor()])
Functions
preprocess async 
preprocess(input: ImageInput) -> Tensor

Preprocess image input.

Parameters:

Name Type Description Default
input ImageInput

Raw image input (bytes, numpy array, PIL Image, or tensor).

 required

Returns:

Type Description
Tensor

Preprocessed tensor ready for model inference.
Source code in inferflow/asyncio/pipeline/classification/torch.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
async def preprocess(self, input: ImageInput) -> torch.Tensor:
    """Preprocess image input.

    Args:
        input: Raw image input (bytes, numpy array, PIL Image, or tensor).

    Returns:
        Preprocessed tensor ready for model inference.
    """

    if isinstance(input, torch.Tensor):
        return self._preprocess_tensor(input)

    image = self._convert_to_pil(input)
    return self._preprocess_image(image)
postprocess async 
postprocess(raw: Tensor) -> ClassificationOutput

Postprocess model output to classification result.

Parameters:

Name Type Description Default
raw Tensor

Raw output tensor from model inference.

 required

Returns:

Type Description
ClassificationOutput

ClassificationOutput with class ID, confidence, and optional class name.
Source code in inferflow/asyncio/pipeline/classification/torch.py 
89
90
91
92
93
94
95
96
97
98
async def postprocess(self, raw: torch.Tensor) -> ClassificationOutput:
    """Postprocess model output to classification result.

    Args:
        raw: Raw output tensor from model inference.

    Returns:
        ClassificationOutput with class ID, confidence, and optional class name.
    """
    return self._postprocess_raw(raw)