| 504 | return out |
| 505 | |
| 506 | class ODE2(nn.Module): |
| 507 | def __init__(self, in_channels,num_slots_N,num_slots_M,iters,resolutions): |
| 508 | super(ODE2, self).__init__() |
| 509 | self.SAn = SlotAttentionModule2(in_channels, resolutions, num_slots_N, iters) |
| 510 | self.SAm = SlotAttentionModule2(in_channels, resolutions, num_slots_M, iters) |
| 511 | self.getalpha = getAlpha(in_channels) |
| 512 | |
| 513 | def forward(self, feature_map): |
| 514 | recon_combined1,f1 = self.SAn(feature_map) |
| 515 | recon_combined2,f2 = self.SAm(f1+feature_map) |
| 516 | alpha = self.getalpha(torch.cat([f1,f2],dim=1)) |
| 517 | out = feature_map+f1*alpha+f2*(1-alpha) |
| 518 | return out,recon_combined1,recon_combined2 |
| 519 | |
| 520 | |
| 521 | class GCM_interFA_ODE_slot(nn.Module): |