(self, x)
| 446 | self.convModules = nn.ModuleList([nn.Conv2d(encoder_dims, encoder_dims//num_slots, kernel_size=1, padding=0, stride=1) for i in range (num_slots)]) |
| 447 | |
| 448 | def forward(self, x): |
| 449 | batch = x.shape[0] |
| 450 | x = einops.rearrange(x, 'b c h w -> b h w c') |
| 451 | x = self.encoder_pos(x) |
| 452 | x = spatial_flatten(x) |
| 453 | x = self.mlp(self.layer_norm(x)) |
| 454 | slots = self.slot_attention(x) |
| 455 | x = spatial_broadcast2(slots, (int(self.resolution), int(self.resolution))) |
| 456 | y = spatial_broadcast(slots, (int(self.resolution), int(self.resolution))) |
| 457 | x = self.decoder_pos(x) |
| 458 | y = self.decoder_pos(y) |
| 459 | z_conv = [] |
| 460 | for i in range(self.num_slots): |
| 461 | z = y[:,i,:,:,:] |
| 462 | z = einops.rearrange(z, 'b h w c -> b c h w') |
| 463 | z1 = self.convModules[i](z) |
| 464 | z_conv.append(z1) |
| 465 | z = torch.stack(z_conv,1) |
| 466 | z = einops.rearrange(z, 'b n c h w -> b (n c) h w') |
| 467 | x = einops.rearrange(x, 'b_n h w c -> b_n c h w') |
| 468 | x = self.decoder_cnn(x) |
| 469 | recons, masks = unstack_and_split(x, batch_size=batch, num_channels=3) |
| 470 | masks = torch.softmax(masks, axis=1) |
| 471 | recon_combined = torch.sum(recons * masks, axis=1) # Recombine image. |
| 472 | return recon_combined, z |
| 473 | |
| 474 | |
| 475 | class getAlpha(nn.Module): |
nothing calls this directly
no test coverage detected