Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

🐞Patch Timm Feature Extractor #714

Merged
merged 3 commits into from
Nov 17, 2022
Merged

🐞Patch Timm Feature Extractor #714

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f480880
Add eval to feture extractor and use torch.no_grad
Nov 16, 2022
83ae8dd
Add eval to forward method
Nov 16, 2022
d937344
Add a requires_grad flag for models that require training the feature…
Nov 16, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions anomalib/models/components/feature_extractors/feature_extractor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,10 @@ class FeatureExtractor(nn.Module):
Args:
backbone (nn.Module): The backbone to which the feature extraction hooks are attached.
layers (Iterable[str]): List of layer names of the backbone to which the hooks are attached.
pre_trained (bool): Whether to use a pre-trained backbone. Defaults to True.
requires_grad (bool): Whether to require gradients for the backbone. Defaults to False.
Models like ``stfpm`` use the feature extractor model as a trainable network. In such cases gradient
computation is required.

Example:
>>> import torch
Expand All @@ -35,11 +39,12 @@ class FeatureExtractor(nn.Module):
[torch.Size([32, 64, 64, 64]), torch.Size([32, 128, 32, 32]), torch.Size([32, 256, 16, 16])]
"""

def __init__(self, backbone: str, layers: List[str], pre_trained: bool = True):
def __init__(self, backbone: str, layers: List[str], pre_trained: bool = True, requires_grad: bool = False):
djdameln marked this conversation as resolved.
Show resolved Hide resolved
super().__init__()
self.backbone = backbone
self.layers = layers
self.idx = self._map_layer_to_idx()
self.requires_grad = requires_grad
self.feature_extractor = timm.create_model(
backbone,
pretrained=pre_trained,
Expand Down Expand Up @@ -85,5 +90,11 @@ def forward(self, input_tensor: Tensor) -> Dict[str, Tensor]:
Returns:
Feature map extracted from the CNN
"""
features = dict(zip(self.layers, self.feature_extractor(input_tensor)))
if self.requires_grad:
features = dict(zip(self.layers, self.feature_extractor(input_tensor)))
else:
self.feature_extractor.eval()
with torch.no_grad():
features = dict(zip(self.layers, self.feature_extractor(input_tensor)))

return features
4 changes: 3 additions & 1 deletion anomalib/models/stfpm/torch_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,9 @@ def __init__(

self.backbone = backbone
self.teacher_model = FeatureExtractor(backbone=self.backbone, pre_trained=True, layers=layers)
self.student_model = FeatureExtractor(backbone=self.backbone, pre_trained=False, layers=layers)
self.student_model = FeatureExtractor(
backbone=self.backbone, pre_trained=False, layers=layers, requires_grad=True
)

# teacher model is fixed
for parameters in self.teacher_model.parameters():
Expand Down