Shortcuts

FeatureExtractor

class mmpretrain.apis.FeatureExtractor(model, pretrained=True, device=None, device_map=None, offload_folder=None, **kwargs)[source]

The inferencer for extract features.

Parameters:
  • model (BaseModel | str | Config) – A model name or a path to the config file, or a BaseModel object. The model name can be found by FeatureExtractor.list_models() and you can also query it in 模型库统计.

  • pretrained (str, optional) – Path to the checkpoint. If None, it will try to find a pre-defined weight from the model you specified (only work if the model is a model name). Defaults to None.

  • device (str, optional) – Device to run inference. If None, the available device will be automatically used. Defaults to None.

  • **kwargs – Other keyword arguments to initialize the model (only work if the model is a model name).

Example

>>> from mmpretrain import FeatureExtractor
>>> inferencer = FeatureExtractor('resnet50_8xb32_in1k', backbone=dict(out_indices=(0, 1, 2, 3)))
>>> feats = inferencer('demo/demo.JPEG', stage='backbone')[0]
>>> for feat in feats:
>>>     print(feat.shape)
torch.Size([256, 56, 56])
torch.Size([512, 28, 28])
torch.Size([1024, 14, 14])
torch.Size([2048, 7, 7])
__call__(inputs, batch_size=1, **kwargs)[source]

Call the inferencer.

Parameters:
  • inputs (str | array | list) – The image path or array, or a list of images.

  • batch_size (int) – Batch size. Defaults to 1.

  • **kwargs – Other keyword arguments accepted by the extract_feat method of the model.

Returns:

The extracted features.

Return type:

tensor | Tuple[tensor]

static list_models(pattern=None)[source]

List all available model names.

Parameters:

pattern (str | None) – A wildcard pattern to match model names.

Returns:

a list of model names.

Return type:

List[str]