Update DL_models.py

DL_models.py

@@ -33,6 +33,10 @@ class CustomResNet(nn.Module):
         self.fc1 = nn.Linear(2 * (512 * 4), 256)
         self.fc2 = nn.Linear(256, 1)
 
+        self.gradients = None
+
+    def activations_hook(self, grad):
+        self.gradients = grad
 
     def forward(self, x1, x2, Location):
         N = x1.shape[0]
@@ -42,7 +46,9 @@ class CustomResNet(nn.Module):
         Location = Location.to(torch.float).to(self.device)
 
         # Process both images through the same ResNet
-        f1 = self.resnet(x1)
+        f1 = self.resnet[:8](x1)
+        h = f1.register_hook(self.activations_hook)
+        f1 = self.resnet[8:](f1)
         f2 = self.resnet(x2)
 
         # Flatten the features
@@ -59,4 +65,13 @@ class CustomResNet(nn.Module):
         x = torch.relu(self.fc1(combined))
         x = self.fc2(x)
 
-        return x
+        return x
+    
+    # method for the gradient extraction
+    def get_activations_gradient(self):
+        return self.gradients
+
+    # method for the activation exctraction
+    def get_activations(self, x):
+        x = self.transform(x.to(torch.float).to(self.device))
+        return self.resnet[:8](x)