rikai.pytorch.models package¶

Submodules¶

rikai.pytorch.models.convnext module¶

rikai.pytorch.models.efficientnet module¶

EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks https://arxiv.org/abs/1905.11946

rikai.pytorch.models.fasterrcnn module¶

rikai.pytorch.models.feature_extractor module¶

class rikai.pytorch.models.feature_extractor.FeatureExtractor(model: Module, node: str, output_field: str = '_rikai_out')¶

Bases: Module

Extract features

eval() → FeatureExtractor¶

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:: Module: self

forward(images: Tensor) → Tensor¶

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool¶

class rikai.pytorch.models.feature_extractor.FeatureExtractorType¶

Bases: TorchModelType

(Experimental) Generic Feature Extractor

Examples

CREATE MODEL resnet_features
FLAVOR pytorch
MODEL_TYPE feature_extractors
OPTIONS (model_type = 'resnet')  # Only resnet is tested
USING '<uri to resnet model>'

SELECT ML_PREDICT(resnet_features, image) AS embedding FROM images

load_model(spec: ModelSpec, **kwargs)¶: Lazy loading the model from a ModelSpec.

predict(*args, **kwargs) → Any¶: Run model inference and convert return types into Rikai-compatible types.

schema() → str¶

Return the string value of model schema.

Examples

>>> model_type.schema()
... "array<struct<box:box2d, score:float, label_id:int>>"

transform() → Callable¶

A callable to pre-process the data before calling inference.

It will be feed into torch.data.DataLoader or tensorflow.data.Dataset.map().

rikai.pytorch.models.keypointrcnn module¶

rikai.pytorch.models.keypointrcnn.label_fn(label_id)¶

rikai.pytorch.models.maskrcnn module¶

rikai.pytorch.models.resnet module¶

ResNet: Deep Residual Learning for Image Recognition https://arxiv.org/pdf/1512.03385.pdf

rikai.pytorch.models.retinanet module¶

rikai.pytorch.models.ssd module¶

rikai.pytorch.models.ssd_class_scores module¶

class rikai.pytorch.models.ssd_class_scores.SSDClassScoresExtractor(backend: SSD, topk_candidates: int = 2)¶

Bases: Module

Extracts the scores (confidences) of each class for all the detected bounding box.

Parameters

backend (torch.nn.Module) – The trained SSD model
topk_candidates (int, optional) – The number of top candidates (classes) returned per box.

Returns

[Dict[str, Tensor]] – With the form of
{"boxes": FloatTensor[N, 4], "labels": Int64Tensor[N, k], "scores": Tensor[N, k]}

forward(images: List[Tensor])¶

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool¶

class rikai.pytorch.models.ssd_class_scores.SSDClassScoresModelType(name: str, pretrained_fn: ~typing.Optional[~typing.Callable] = None, label_fn: ~typing.Optional[~typing.Callable[[int], str]] = <function detection_label_fn>, collate_fn: ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]] = <function detection_collate_fn>, register: bool = True)¶

Bases: ObjectDetectionModelType

DEFAULT_MIN_SCORE = 0.3¶

load_model(spec: ModelSpec, **kwargs)¶: Lazy loading the model from a ModelSpec.

predict(images, *args, **kwargs) → List¶: Run model inference and convert return types into Rikai-compatible types.

schema() → str¶

Return the string value of model schema.

Examples

>>> model_type.schema()
... "array<struct<box:box2d, score:float, label_id:int>>"

rikai.pytorch.models.torch module¶

ModelType for official torchvision models

class rikai.pytorch.models.torch.ClassificationModelType(name: str, pretrained_fn: ~typing.Optional[~typing.Callable] = None, label_fn: ~typing.Optional[~typing.Callable[[int], str]] = <function classification_label_fn>, register: bool = True)¶

Bases: TorchModelType

Shared ModelType for image classification

predict(images, *args, **kwargs) → Any¶: Run model inference and convert return types into Rikai-compatible types.

schema() → str¶

Return the string value of model schema.

Examples

>>> model_type.schema()
... "array<struct<box:box2d, score:float, label_id:int>>"

transform() → Callable¶

A callable to pre-process the data before calling inference.

It will be feed into torch.data.DataLoader or tensorflow.data.Dataset.map().

class rikai.pytorch.models.torch.ObjectDetectionModelType(name: str, pretrained_fn: ~typing.Optional[~typing.Callable] = None, label_fn: ~typing.Optional[~typing.Callable[[int], str]] = <function detection_label_fn>, collate_fn: ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]] = <function detection_collate_fn>, register: bool = True)¶

Bases: TorchModelType

Shared ModelType for object detections in Torchvision

https://pytorch.org/vision/stable/models.html

predict(images, *args, **kwargs) → Any¶: Run model inference and convert return types into Rikai-compatible types.

schema() → str¶

Return the string value of model schema.

Examples

>>> model_type.schema()
... "array<struct<box:box2d, score:float, label_id:int>>"

transform() → Callable¶

A callable to pre-process the data before calling inference.

It will be feed into torch.data.DataLoader or tensorflow.data.Dataset.map().

class rikai.pytorch.models.torch.TorchModelType(name: str, pretrained_fn: Optional[Callable] = None, label_fn: Optional[Callable[[int], str]] = None, collate_fn: Optional[Callable] = None, register: bool = True)¶

Bases: ModelType, Pretrained, ABC

Base ModelType for PyTorch models.

load_model(spec: ModelSpec, **kwargs)¶: Lazy loading the model from a ModelSpec.

pretrained_model() → Any¶

release()¶

Release underneath resources if applicable.

It will be called after a model runner finishes a partition in Spark.

Module contents¶

Rikai-implemented PyTorch models and executors.