-
Notifications
You must be signed in to change notification settings - Fork 23
feature_extraction
antonemanuel edited this page Aug 20, 2021
·
1 revision
Provides classes and functions for extracting embedding vectors from neural networks.
class ResnetEmbeddingsExtractor(torch.nn.Module)A class to hold, and extract embedding vectors from, a resnet.
Attributes:
-
backbone- The resnet from which to extract embedding vectors.
| __init__(backbone_name: str, device: torch.device) -> NoneConstruct the backbone and set appropriate mode and device
Arguments:
-
backbone_name- The name of the desired backbone. Must be one of: [resnet18, wide_resnet50]. -
device- The device where to run the network.
| to_device(device: torch.device) -> NonePerform device conversion on backone
See pytorch docs for documentation on torch.Tensor.to
| forward(batch: torch.Tensor, channel_indices: Optional[torch.Tensor] = None, layer_hook: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, layer_indices: Optional[List[int]] = None) -> torch.TensorRun inference on backbone and return the embedding vectors.
Arguments:
-
batch- A batch of images. -
channel_indices- A list of indices with the desired channels to include in the embedding vectors. -
layer_hook- A function that runs on each layer of the resnet before concatenating them. -
layer_indices- A list of indices with the desired layers to include in the embedding vectors.
Returns:
-
embedding_vectors- The embedding vectors.
| from_dataloader(dataloader: DataLoader, channel_indices: Optional[torch.Tensor] = None, layer_hook: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, layer_indices: Optional[List[int]] = None) -> torch.TensorSame as self.forward but take a dataloader instead of a tensor as argument.
concatenate_layers(layers: List[torch.Tensor]) -> torch.TensorScale all tensors to the heigth and width of the first tensor and concatenate them.
concatenate_two_layers(layer1: torch.Tensor, layer2: torch.Tensor) -> torch.TensorScale the second tensor to the height and width of the first tensor and concatenate them.