(self, encoder_dims, resolution, num_slots, iters)
| 378 | |
| 379 | class SlotAttentionModule(nn.Module): |
| 380 | def __init__(self, encoder_dims, resolution, num_slots, iters): |
| 381 | super(SlotAttentionModule, self).__init__() |
| 382 | self.resolution = resolution |
| 383 | self.encoder_pos = SoftPositionEmbed(encoder_dims, ((int(resolution), int(resolution)))) |
| 384 | self.layer_norm = nn.LayerNorm(encoder_dims) |
| 385 | self.mlp = nn.Sequential( |
| 386 | nn.Linear(encoder_dims, encoder_dims), |
| 387 | nn.ReLU(inplace=True), |
| 388 | nn.Linear(encoder_dims, encoder_dims) |
| 389 | ) |
| 390 | self.slot_attention = SlotAttention(iters=iters, |
| 391 | num_slots=num_slots, |
| 392 | encoder_dims=encoder_dims, |
| 393 | hidden_dim=encoder_dims) |
| 394 | self.decoder_pos = SoftPositionEmbed(encoder_dims, (int(resolution), int(resolution))) |
| 395 | self.conv = nn.Conv2d(encoder_dims*num_slots, encoder_dims, kernel_size=1, padding=0, stride=1) |
| 396 | |
| 397 | def forward(self, x): |
| 398 | x = einops.rearrange(x, 'b c h w -> b h w c') |
nothing calls this directly
no test coverage detected