feat: bilinear interpolation of attention heads when dimension does n…

…ot match, useful for segformer G

models/modules/attn_network.py

@@ -1,4 +1,6 @@
 from torch import nn
+import torch.nn.functional as F
+import warnings
 
 
 class BaseGenerator_attn(nn.Module):
@@ -12,9 +14,24 @@ def compute_outputs(self, input, attentions, images):
         outputs = []
 
         for i in range(self.nb_mask_attn - self.nb_mask_input):
-            outputs.append(images[i] * attentions[i])
+            if images[i].shape == attentions[i].shape:
+                outputs.append(images[i] * attentions[i])
+            else:
+                warnings.warn("Bilinear interpolation of attention heads")
+                rattention = F.interpolate(
+                    attentions[i], size=(images[i].shape[2], images[i].shape[3])
+                )
+                outputs.append(images[i] * rattention)
+
         for i in range(self.nb_mask_attn - self.nb_mask_input, self.nb_mask_attn):
-            outputs.append(input * attentions[i])
+            if input.shape == attentions[i]:
+                outputs.append(input * attentions[i])
+            else:
+                warnings.warn("Bilinear interpolation of attention heads")
+                rattention = F.interpolate(
+                    attentions[i], size=(input.shape[2], input.shape[3])
+                )
+                outputs.append(input * rattention)
 
         return images, attentions, outputs