upload

alignedthreeattn_model.py

@@ -54,7 +54,12 @@ class ThreeAttnNodes(nn.Module):
         x = F.interpolate(x, size=(588, 588), mode="bilinear")
         feat2 = self.backbone2(x)
         feats = torch.cat([feat1, feat2, feat3], dim=1)
-        out = torch.einsum("b l p i, l o i -> b l p o", feats, self.align_weights)
+        # out = torch.einsum("b l p i, l o i -> b l p o", feats, self.align_weights)
+        outs = []
+        for i_layer in range(36):
+            out = torch.einsum("b p i, o i -> b p o", feats[:, i_layer], self.align_weights[i_layer])
+            outs.append(out)
+        out = torch.stack(outs, dim=1)
         out = rearrange(out[:, :, 1:], "b l (h w) o -> b l h w o", h=42, w=42)
         return out