Method forward

lib/Slot.py:231–263

(self, image)

    def forward(self, image):
        # `image` has shape: [batch_size, num_channels, height, width].
        # Convolutional encoder with position embedding.
        x = self.encoder_cnn(image)  # CNN Backbone.
        x = einops.rearrange(x, 'b c h w -> b h w c')
        x = self.encoder_pos(x)  # Position embedding.
        x = spatial_flatten(x)  # Flatten spatial dimensions (treat image as set).
        x = self.mlp(self.layer_norm(x))  # Feedforward network on set.
        # `x` has shape: [batch_size, width*height, input_size].

        # Slot Attention module.
        slots = self.slot_attention(x)
        # `slots` has shape: [batch_size, num_slots, slot_size].

        # Spatial broadcast decoder.
        x = spatial_broadcast(slots, self.decoder_initial_size)
        # `x` has shape: [batch_size*num_slots, height_init, width_init, slot_size].
        x = self.decoder_pos(x)
        x = einops.rearrange(x, 'b_n h w c -> b_n c h w')
        x = self.decoder_cnn(x)
        # `x` has shape: [batch_size*num_slots, num_channels+1, height, width].
        
        # Undo combination of slot and batch dimension; split alpha masks.
        recons, masks = unstack_and_split(x, batch_size=image.shape[0], num_channels=self.in_out_channels)
        # `recons` has shape: [batch_size, num_slots, num_channels, height, width].
        # `masks` has shape: [batch_size, num_slots, 1, height, width].
        
        # Normalize alpha masks over slots.
        masks = torch.softmax(masks, axis=1)

        recon_combined = torch.sum(recons * masks, axis=1)  # Recombine image.
        # `recon_combined` has shape: [batch_size, num_channels, height, width].
        return recon_combined, recons, masks, slots

if __name__ =='__main__':
    x = torch.rand(4,128,12,12).cuda()