Fixed bugs

Signed-off-by: Devansh Agarwal <[email protected]>

Author: code-dev05

src/anomalib/models/components/feature_extractors/network_feature_extractor.py

@@ -29,7 +29,6 @@ def __init__(self, backbone, layers_to_extract_from, pre_trained=False):
         for extract_layer in layers_to_extract_from:
             self.register_hook(extract_layer)
 
-        self.to(self.device)
 
     def forward(self, images, eval=True):
         self.outputs.clear()
@@ -45,7 +44,7 @@ def forward(self, images, eval=True):
 
     def feature_dimensions(self, input_shape):
         """Computes the feature dimensions for all layers given input_shape."""
-        _input = torch.ones([1] + list(input_shape)).to(self.device)
+        _input = torch.ones([1] + list(input_shape))
         _output = self(_input)
         return [_output[layer].shape[1] for layer in self.layers_to_extract_from]
 

src/anomalib/models/image/glass/lightning_model.py

@@ -16,9 +16,10 @@
 from anomalib.visualization import Visualizer
 
 from .loss import FocalLoss
-from .perlin import PerlinNoise
 from .torch_model import GlassModel
 
+from anomalib.data.utils.generators.perlin import PerlinAnomalyGenerator
+
 
 class Glass(AnomalibModule):
     def __init__(
@@ -54,7 +55,7 @@ def __init__(
             visualizer=visualizer,
         )
 
-        self.perlin = PerlinNoise(anomaly_source_path)
+        self.augmentor = PerlinAnomalyGenerator(anomaly_source_path)
 
         self.model = GlassModel(
             input_shape=input_shape,
@@ -82,45 +83,41 @@ def __init__(
 
         self.focal_loss = FocalLoss()
 
-    def configure_optimizers(self) -> list[optim.Optimizer]:
-        optimizers = []
-        if not self.model.pre_trained:
-            backbone_opt = optim.AdamW(self.model.foward_modules["feature_aggregator"].backbone.parameters(), self.lr)
-            optimizers.append(backbone_opt)
+        if pre_proj > 0:
+            self.proj_opt = optim.AdamW(self.model.pre_projection.parameters(), self.lr, weight_decay=1e-5)
         else:
-            optimizers.append(None)
+            self.proj_opt = None
 
-        if self.model.pre_proj > 0:
-            proj_opt = optim.AdamW(self.model.pre_projection.parameters(), self.lr, weight_decay=1e-5)
-            optimizers.append(proj_opt)
+        if not pre_trained:
+            self.backbone_opt = optim.AdamW(self.model.foward_modules["feature_aggregator"].backbone.parameters(), self.lr)
         else:
-            optimizers.append(None)
+            self.backbone_opt = None
 
+    def configure_optimizers(self) -> list[optim.Optimizer]:
         dsc_opt = optim.AdamW(self.model.discriminator.parameters(), lr=self.lr * 2)
-        optimizers.append(dsc_opt)
 
-        return optimizers
+        return dsc_opt
 
     def training_step(
         self,
         batch: Batch,
         batch_idx: int,
     ) -> STEP_OUTPUT:
-        backbone_opt, proj_opt, dsc_opt = self.optimizers()
+        dsc_opt = self.optimizers()
 
         self.model.forward_modules.eval()
         if self.model.pre_proj > 0:
             self.pre_projection.train()
         self.model.discriminator.train()
 
         dsc_opt.zero_grad()
-        if proj_opt is not None:
-            proj_opt.zero_grad()
-        if backbone_opt is not None:
-            backbone_opt.zero_grad()
+        if self.proj_opt is not None:
+            self.proj_opt.zero_grad()
+        if self.backbone_opt is not None:
+            self.backbone_opt.zero_grad()
 
         img = batch.image
-        aug, mask_s = self.perlin(img)
+        aug, mask_s = self.augmentor(img)
 
         true_feats, fake_feats = self.model(img, aug)
 
@@ -191,10 +188,10 @@ def training_step(
         loss = bce_loss + focal_loss
         loss.backward()
 
-        if proj_opt is not None:
-            proj_opt.step()
-        if backbone_opt is not None:
-            backbone_opt.step()
+        if self.proj_opt is not None:
+            self.proj_opt.step()
+        if self.backbone_opt is not None:
+            self.backbone_opt.step()
         dsc_opt.step()
 
     def on_train_start(self) -> None:

src/anomalib/models/image/glass/torch_model.py

@@ -198,7 +198,7 @@ def __init__(
         self.layers = layers
         self.input_shape = input_shape
 
-        self.forward_modules = torch.ModuleDict({})
+        self.forward_modules = torch.nn.ModuleDict({})
         feature_aggregator = NetworkFeatureAggregator(
             self.backbone,
             self.layers,
@@ -257,7 +257,7 @@ def generate_embeddings(self, images, provide_patch_shapes=False, eval=False):
             with torch.no_grad():
                 features = self.forward_modules["feature_aggregator"](images)
 
-        features = [features[layer] for layer in self.layers_to_extract_from]
+        features = [features[layer] for layer in self.layers]
         for i, feat in enumerate(features):
             if len(feat.shape) == 3:
                 B, L, C = feat.shape