Skip to content

ocl.feature_extractors

Implementation of feature extractors that can be used for object centric learning.

These are grouped into 3 modules

Utilities used by all modules are found in ocl.feature_extractors.utils.

Important note: In order to use feature extractors in timm and clip this package has to be installed with the timm and/or clip extras (see Installation for further information on installing extras).

SlotAttentionFeatureExtractor

Bases: ImageFeatureExtractor

CNN-based feature extractor as used in the slot attention paper.

Reference: Locatello et al., Object-Centric Learning with Slot Attention, NeurIPS 2020

Source code in ocl/feature_extractors/misc.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class SlotAttentionFeatureExtractor(ImageFeatureExtractor):
    """CNN-based feature extractor as used in the slot attention paper.

    Reference: Locatello et al., Object-Centric Learning with Slot Attention, NeurIPS 2020
    """

    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
        )

    @property
    def feature_dim(self):
        return 64

    def forward_images(self, images: torch.Tensor):
        features = self.layers(images)
        flattened, positions = cnn_compute_positions_and_flatten(features)
        return flattened, positions

SAViFeatureExtractor

Bases: ImageFeatureExtractor

CNN-based feature extractor as used in the slot attention for video paper.

Reference: Kipf et al., Conditional Object-Centric Learning from Video, ICLR 2020

Source code in ocl/feature_extractors/misc.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
class SAViFeatureExtractor(ImageFeatureExtractor):
    """CNN-based feature extractor as used in the slot attention for video paper.

    Reference: Kipf et al., Conditional Object-Centric Learning from Video, ICLR 2020
    """

    def __init__(self, larger_input_arch: bool = False):
        """Initialize SAVi feature extractor.

        Args:
            larger_input_arch: Use the architecture for larger image datasets such as MOVi++, which
                contains more a stride in the first layer and a higher number of feature channels in
                the CNN backbone.
        """
        super().__init__()
        self.larger_input_arch = larger_input_arch
        if larger_input_arch:
            self.layers = nn.Sequential(
                # Pytorch does not support stride>1 with padding=same.
                # Implement tensorflow behaviour manually.
                # See: https://discuss.pytorch.org/t/same-padding-equivalent-in-pytorch/85121/4
                nn.ZeroPad2d((1, 2, 1, 2)),
                nn.Conv2d(3, out_channels=64, kernel_size=5, stride=2, padding="valid"),
                nn.ReLU(inplace=True),
                nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            )
        else:
            self.layers = nn.Sequential(
                nn.Conv2d(3, out_channels=32, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(32, out_channels=32, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(32, out_channels=32, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(32, out_channels=32, kernel_size=5, padding="same"),
            )

    @property
    def feature_dim(self):
        return 64 if self.larger_input_arch else 32

    def forward_images(self, images: torch.Tensor):
        features = self.layers(images)
        flattened, positions = cnn_compute_positions_and_flatten(features)
        return flattened, positions

__init__

Initialize SAVi feature extractor.

Parameters:

Name Type Description Default
larger_input_arch bool

Use the architecture for larger image datasets such as MOVi++, which contains more a stride in the first layer and a higher number of feature channels in the CNN backbone.

 False
Source code in ocl/feature_extractors/misc.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def __init__(self, larger_input_arch: bool = False):
    """Initialize SAVi feature extractor.

    Args:
        larger_input_arch: Use the architecture for larger image datasets such as MOVi++, which
            contains more a stride in the first layer and a higher number of feature channels in
            the CNN backbone.
    """
    super().__init__()
    self.larger_input_arch = larger_input_arch
    if larger_input_arch:
        self.layers = nn.Sequential(
            # Pytorch does not support stride>1 with padding=same.
            # Implement tensorflow behaviour manually.
            # See: https://discuss.pytorch.org/t/same-padding-equivalent-in-pytorch/85121/4
            nn.ZeroPad2d((1, 2, 1, 2)),
            nn.Conv2d(3, out_channels=64, kernel_size=5, stride=2, padding="valid"),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
        )
    else:
        self.layers = nn.Sequential(
            nn.Conv2d(3, out_channels=32, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, out_channels=32, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, out_channels=32, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, out_channels=32, kernel_size=5, padding="same"),
        )

DVAEFeatureExtractor

Bases: ImageFeatureExtractor

DVAE VQ Encoder as used in SLATE.

Reference

Singh et al., Simple Unsupervised Object-Centric Learning for Complex and Naturalistic Videos, NeurIPS 2022

Source code in ocl/feature_extractors/misc.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
class DVAEFeatureExtractor(ImageFeatureExtractor):
    """DVAE VQ Encoder as used in SLATE.

    Reference:
        Singh et al., Simple Unsupervised Object-Centric Learning for Complex and
        Naturalistic Videos, NeurIPS 2022
    """

    def __init__(
        self,
        encoder: nn.Module,
        positional_encoder: nn.Module,
        dictionary: nn.Module,
        tau: float = 1.0,
        hard: bool = False,
    ):
        """Feature extractor as used in the SLATE paper.

        Args:
            encoder: torch Module that transforms image to the patch representations.
            positional_encoder: torch Module that adds pos encoding.
            dictionary: map from onehot vectors to embeddings.
            tau: temporature for gumbel_softmax.
            hard: hard gumbel_softmax if True.
        """
        super().__init__()
        self.tau = tau
        self.hard = hard
        self.dictionary = dictionary
        self.positional_encoder = positional_encoder
        self.encoder = encoder

    @property
    def feature_dim(self):
        return 64

    def forward_images(self, images: torch.Tensor):
        z_logits = nn.functional.log_softmax(self.encoder(images), dim=1)
        _, _, H_enc, W_enc = z_logits.size()
        z = nn.functional.gumbel_softmax(z_logits, float(self.tau), self.hard, dim=1)
        z_hard = nn.functional.gumbel_softmax(z_logits, float(self.tau), True, dim=1).detach()

        # add beginning of sequence (BOS) token
        # [1, 0, 0, 0, ...] is encoding for BOS token
        # and each sequence starts from such token
        z_hard = z_hard.permute(0, 2, 3, 1).flatten(start_dim=1, end_dim=2)
        # add first zeros column to the z_hard matrix
        z_transformer_input = torch.cat([torch.zeros_like(z_hard[..., :1]), z_hard], dim=-1)
        # add first zeros row to the z_hard matrix
        z_transformer_input = torch.cat(
            [torch.zeros_like(z_transformer_input[..., :1, :]), z_transformer_input], dim=-2
        )
        # fill new row and column with one,
        # so that we added [1, 0, 0, 0, ...] token
        z_transformer_input[:, 0, 0] = 1.0

        # tokens to embeddings
        features = self.dictionary(z_transformer_input)
        features = self.positional_encoder(features)

        slot_attention_features = features[:, 1:]

        transformer_input = features[:, :-1]
        aux_features = {
            "z": z,
            "targets": transformer_input,
            "z_hard": z_hard,
        }
        return slot_attention_features, None, aux_features

__init__

Feature extractor as used in the SLATE paper.

Parameters:

Name Type Description Default
encoder nn.Module

torch Module that transforms image to the patch representations.

 required
positional_encoder nn.Module

torch Module that adds pos encoding.

 required
dictionary nn.Module

map from onehot vectors to embeddings.

 required
tau float

temporature for gumbel_softmax.

 1.0
hard bool

hard gumbel_softmax if True.

 False
Source code in ocl/feature_extractors/misc.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def __init__(
    self,
    encoder: nn.Module,
    positional_encoder: nn.Module,
    dictionary: nn.Module,
    tau: float = 1.0,
    hard: bool = False,
):
    """Feature extractor as used in the SLATE paper.

    Args:
        encoder: torch Module that transforms image to the patch representations.
        positional_encoder: torch Module that adds pos encoding.
        dictionary: map from onehot vectors to embeddings.
        tau: temporature for gumbel_softmax.
        hard: hard gumbel_softmax if True.
    """
    super().__init__()
    self.tau = tau
    self.hard = hard
    self.dictionary = dictionary
    self.positional_encoder = positional_encoder
    self.encoder = encoder