| 139 | |
| 140 | |
| 141 | class MemoryModulev2(nn.Module): |
| 142 | def __init__(self, dim, num_head=8, window=7, norm_cfg=dict(type="SyncBN", requires_grad=True)): |
| 143 | super().__init__() |
| 144 | self.num_head = num_head |
| 145 | self.window = window |
| 146 | |
| 147 | self.q = nn.Linear(dim, dim) |
| 148 | self.a = nn.Linear(dim, dim) |
| 149 | self.l = nn.Linear(dim, dim) |
| 150 | self.conv = nn.Conv2d(dim, dim, 7, padding=3, groups=dim) |
| 151 | |
| 152 | self.proj = nn.Linear(dim, dim) |
| 153 | if window != 0: |
| 154 | self.t = nn.Linear(dim, window * window * num_head) |
| 155 | self.conv_att = nn.Conv2d(dim, dim, 7, padding=3, groups=dim) |
| 156 | self.kv = nn.Linear(dim, dim * 2) |
| 157 | self.m = nn.Parameter(torch.zeros(1, window, window, dim), requires_grad=True) |
| 158 | self.proj = nn.Linear(dim * 2, dim) |
| 159 | |
| 160 | self.act = nn.GELU() |
| 161 | self.norm = LayerNorm(dim, eps=1e-6, data_format="channels_last") |
| 162 | |
| 163 | def forward(self, x): |
| 164 | B, H, W, C = x.size() |
| 165 | x = self.norm(x) |
| 166 | |
| 167 | q = self.q(x) |
| 168 | x = self.l(x).permute(0, 3, 1, 2) |
| 169 | x = self.act(x) |
| 170 | |
| 171 | a = self.conv(x) |
| 172 | a = a.permute(0, 2, 3, 1) |
| 173 | a = self.a(a) |
| 174 | |
| 175 | if self.window != 0: |
| 176 | b = self.conv_att(x) |
| 177 | b = b.permute(0, 2, 3, 1) |
| 178 | kv = self.kv(b) |
| 179 | kv = kv.reshape(B, H * W, 2, self.num_head, C // self.num_head).permute(2, 0, 3, 1, 4) |
| 180 | k, v = kv.unbind(0) |
| 181 | t = self.t(a).reshape(B, H * W, self.num_head, -1).permute(0, 2, 1, 3).softmax(dim=-1) |
| 182 | |
| 183 | m = self.m.reshape(1, -1, self.num_head, C // self.num_head).permute(0, 2, 1, 3).expand(B, -1, -1, -1) |
| 184 | attn = (m * (C // self.num_head) ** -0.5) @ k.transpose(-2, -1) |
| 185 | attn = attn.softmax(dim=-1) |
| 186 | attn = (attn @ v).reshape(B, self.num_head, self.window * self.window, C // self.num_head) |
| 187 | # attn = F.interpolate(attn, (H, W), mode='bilinear', align_corners=False).permute(0, 2, 3, 1) |
| 188 | attn = (t @ attn).permute(0, 2, 1, 3).reshape(B, H, W, C) |
| 189 | |
| 190 | x = q * a |
| 191 | |
| 192 | if self.window != 0: |
| 193 | x = torch.cat([x, attn], dim=3) |
| 194 | x = self.proj(x) |
| 195 | |
| 196 | return x |
| 197 | |
| 198 | |