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.
- 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
- predict(*args, **kwargs) Any ¶
Run model inference and convert return types into Rikai-compatible types.
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 modeltopk_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.
- 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¶
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.
- 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.
- 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.
- 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.