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 byFeatureExtractor.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])