| 490 | |
| 491 | |
| 492 | class ODE(nn.Module): |
| 493 | def __init__(self, in_channels,num_slots_N,num_slots_M,iters,resolutions): |
| 494 | super(ODE, self).__init__() |
| 495 | self.SAn = SlotAttentionModule(in_channels, resolutions, num_slots_N, iters) |
| 496 | self.SAm = SlotAttentionModule(in_channels, resolutions, num_slots_M, iters) |
| 497 | self.getalpha = getAlpha(in_channels) |
| 498 | |
| 499 | def forward(self, feature_map): |
| 500 | f1 = self.SAn(feature_map) |
| 501 | f2 = self.SAm(f1+feature_map) |
| 502 | alpha = self.getalpha(torch.cat([f1,f2],dim=1)) |
| 503 | out = feature_map+f1*alpha+f2*(1-alpha) |
| 504 | return out |
| 505 | |
| 506 | class ODE2(nn.Module): |
| 507 | def __init__(self, in_channels,num_slots_N,num_slots_M,iters,resolutions): |