| 116 | |
| 117 | |
| 118 | class CrossPath(nn.Module): |
| 119 | def __init__(self, dim, reduction=1, num_heads=None, norm_layer=nn.LayerNorm): |
| 120 | super().__init__() |
| 121 | self.channel_proj1 = nn.Linear(dim, dim // reduction * 2) |
| 122 | self.channel_proj2 = nn.Linear(dim, dim // reduction * 2) |
| 123 | self.act1 = nn.ReLU(inplace=True) |
| 124 | self.act2 = nn.ReLU(inplace=True) |
| 125 | self.cross_attn = CrossAttention(dim // reduction, num_heads=num_heads) |
| 126 | self.end_proj1 = nn.Linear(dim // reduction * 2, dim) |
| 127 | self.end_proj2 = nn.Linear(dim // reduction * 2, dim) |
| 128 | self.norm1 = norm_layer(dim) |
| 129 | self.norm2 = norm_layer(dim) |
| 130 | |
| 131 | def forward(self, x1, x2): |
| 132 | y1, u1 = self.act1(self.channel_proj1(x1)).chunk(2, dim=-1) |
| 133 | y2, u2 = self.act2(self.channel_proj2(x2)).chunk(2, dim=-1) |
| 134 | v1, v2 = self.cross_attn(u1, u2) |
| 135 | y1 = torch.cat((y1, v1), dim=-1) |
| 136 | y2 = torch.cat((y2, v2), dim=-1) |
| 137 | out_x1 = self.norm1(x1 + self.end_proj1(y1)) |
| 138 | out_x2 = self.norm2(x2 + self.end_proj2(y2)) |
| 139 | return out_x1, out_x2 |
| 140 | |
| 141 | |
| 142 | # Stage 2 |