class_docstr: Extract features from a CNN.
  
  Args:
      backbone (str | BackboneParams | dict | nn.Module): The backbone to which the feature extraction hooks are
          attached. If the name is provided, the model is loaded from torchvision. Otherwise, the model class can be
          provided and it will try to load the weights from the provided weights file. Last, an instance of nn.Module
          can also be passed directly.
      return_nodes (Iterable[str]): List of layer names of the backbone to which the hooks are attached.
          You can find the names of these nodes by using ``get_graph_node_names`` function.
      weights (str | WeightsEnum | None): Weights enum to use for the model. Torchvision models require
          ``WeightsEnum``. These enums are defined in ``torchvision.models.<model>``. You can pass the weights
          path for custom models.
      requires_grad (bool): Models like ``stfpm`` use the feature extractor for training. In such cases we should
          set ``requires_grad`` to ``True``. Default is ``False``.
      tracer_kwargs (dict | None): a dictionary of keyword arguments for NodePathTracer (which passes them onto
          it's parent class torch.fx.Tracer). Can be used to allow not tracing through a list of problematic
          modules, by passing a list of `leaf_modules` as one of the `tracer_kwargs`.
  
  Example:
      With torchvision models:
  
          >>> import torch
          >>> from anomalib.models.components.feature_extractors import TorchFXFeatureExtractor
          >>> from torchvision.models.efficientnet import EfficientNet_B5_Weights
          >>> feature_extractor = TorchFXFeatureExtractor(
                  backbone="efficientnet_b5",
                  return_nodes=["features.6.8"],
                  weights=EfficientNet_B5_Weights.DEFAULT
              )
          >>> input = torch.rand((32, 3, 256, 256))
          >>> features = feature_extractor(input)
          >>> [layer for layer in features.keys()]
              ["features.6.8"]
          >>> [feature.shape for feature in features.values()]
              [torch.Size([32, 304, 8, 8])]
  
      With custom models:
  
          >>> import torch
          >>> from anomalib.models.components.feature_extractors import TorchFXFeatureExtractor
          >>> feature_extractor = TorchFXFeatureExtractor(
                  "path.to.CustomModel", ["linear_relu_stack.3"], weights="path/to/weights.pth"
              )
          >>> input = torch.randn(1, 1, 28, 28)
          >>> features = feature_extractor(input)
          >>> [layer for layer in features.keys()]
              ["linear_relu_stack.3"]
  
      with model instances:
  
          >>> import torch
          >>> from anomalib.models.components.feature_extractors import TorchFXFeatureExtractor
          >>> from timm import create_model
          >>> model = create_model("resnet18", pretrained=True)
          >>> feature_extractor = TorchFXFeatureExtractor(model, ["layer1"])
          >>> input = torch.rand((32, 3, 256, 256))
          >>> features = feature_extractor(input)
          >>> [layer for layer in features.keys()]
              ["layer1"]
          >>> [feature.shape for feature in features.values()]
              [torch.Size([32, 64, 64, 64])]
class_docstr_tok: ["Extract","features","from","a","CNN",".","Args",":","backbone","(","str","|","BackboneParams","|","dict","|","nn.Module",")",":","The","backbone","to","which","the","feature","extraction","hooks","are","attached",".","If","the","name","is","provided",",","the","model","is","loaded","from","torchvision",".","Otherwise",",","the","model","class","can","be","provided","and","it","will","try","to","load","the","weights","from","the","provided","weights","file",".","Last",",","an","instance","of","nn.Module","can","also","be","passed","directly",".","return_nodes","(","Iterable","[","str","]",")",":","List","of","layer","names","of","the","backbone","to","which","the","hooks","are","attached",".","You","can","find","the","names","of","these","nodes","by","using","``","get_graph_node_names","``","function",".","weights","(","str","|","WeightsEnum","|","None",")",":","Weights","enum","to","use","for","the","model",".","Torchvision","models","require","``","WeightsEnum","``",".","These","enums","are","defined","in","``","torchvision.models.","<","model",">","``",".","You","can","pass","the","weights","path","for","custom","models",".","requires_grad","(","bool",")",":","Models","like","``","stfpm","``","use","the","feature","extractor","for","training",".","In","such","cases","we","should","set","``","requires_grad","``","to","``","True","``",".","Default","is","``","False","``",".","tracer_kwargs","(","dict","|","None",")",":","a","dictionary","of","keyword","arguments","for","NodePathTracer","(","which","passes","them","onto","it","'s","parent","class","torch.fx.Tracer",")",".","Can","be","used","to","allow","not","tracing","through","a","list","of","problematic","modules",",","by","passing","a","list","of","`","leaf_modules","`","as","one","of","the","`","tracer_kwargs","`",".","Example",":","With","torchvision","models",":",">",">",">","import","torch",">",">",">","from","anomalib.models.components.feature_extractors","import","TorchFXFeatureExtractor",">",">",">","from","torchvision.models.efficientnet","import","EfficientNet_B5_Weights",">",">",">","feature_extractor","=","TorchFXFeatureExtractor","(","backbone=","''","efficientnet_b5","''",",","return_nodes=","[","``","features.6.8","''","]",",","weights=EfficientNet_B5_Weights.DEFAULT",")",">",">",">","input","=","torch.rand","(","(","32",",","3",",","256",",","256",")",")",">",">",">","features","=","feature_extractor","(","input",")",">",">",">","[","layer","for","layer","in","features.keys","(",")","]","[","``","features.6.8","''","]",">",">",">","[","feature.shape","for","feature","in","features.values","(",")","]","[","torch.Size","(","[","32",",","304",",","8",",","8","]",")","]","With","custom","models",":",">",">",">","import","torch",">",">",">","from","anomalib.models.components.feature_extractors","import","TorchFXFeatureExtractor",">",">",">","feature_extractor","=","TorchFXFeatureExtractor","(","``","path.to.CustomModel","''",",","[","``","linear_relu_stack.3","''","]",",","weights=","''","path\/to\/weights.pth","''",")",">",">",">","input","=","torch.randn","(","1",",","1",",","28",",","28",")",">",">",">","features","=","feature_extractor","(","input",")",">",">",">","[","layer","for","layer","in","features.keys","(",")","]","[","``","linear_relu_stack.3","''","]","with","model","instances",":",">",">",">","import","torch",">",">",">","from","anomalib.models.components.feature_extractors","import","TorchFXFeatureExtractor",">",">",">","from","timm","import","create_model",">",">",">","model","=","create_model","(","``","resnet18","''",",","pretrained=True",")",">",">",">","feature_extractor","=","TorchFXFeatureExtractor","(","model",",","[","``","layer1","''","]",")",">",">",">","input","=","torch.rand","(","(","32",",","3",",","256",",","256",")",")",">",">",">","features","=","feature_extractor","(","input",")",">",">",">","[","layer","for","layer","in","features.keys","(",")","]","[","``","layer1","''","]",">",">",">","[","feature.shape","for","feature","in","features.values","(",")","]","[","torch.Size","(","[","32",",","64",",","64",",","64","]",")","]"]
docstr: Extract features from the input.
docstr_tok: ["Extract","features","from","the","input","."]
code: def forward(self, inputs: Tensor) -> dict[str, Tensor]:
      """Extract features from the input."""
      return self.feature_extractor(inputs)
  
file_path: anomalib/src/anomalib/models/components/feature_extractors/torchfx.py
imports: from __future__ import annotations
  import importlib
  from dataclasses import dataclass, field
  from typing import Callable
  import torch
  from torch import Tensor, nn
  from torch.fx.graph_module import GraphModule
  from torchvision.models._api import WeightsEnum
  from torchvision.models.feature_extraction import create_feature_extractor
label_desc: Use image node_id 4 for calling the TorchFXFeatureExtractor obj's underlying member method code with example usage: obj.forward(inputs) and returns: self