[FEATURE] Method to convert feature embeddings into predictions

[crypdick]

Is your feature request related to a problem? Please describe.
Currently, saving both the feature embedding vector and the prediction vector requires two forward passes:

predictions = model(inputs)
embeddings = model.forward_features(inputs)  # inefficient

This is technically not necessary, since the embeddings are computed when you compute predictions.

In addition, once you have your embeddings, there is no standardized method to convert the embeddings into predictions. For example, with ResNet Models I could do something like:

predictions = model.fc(embeddings)

...but that does not generalize to other models since not every model has a fc layer.

Describe the solution you'd like

For every model to have a forward method (I suggest the name forward_predictions) which takes an embedding as input and outputs a prediction.

For example, ResNet would go from:

def forward(self, x):
    x = self.forward_features(x)
    x = self.global_pool(x)
    if self.drop_rate:
        x = F.dropout(x, p=float(self.drop_rate), training=self.training)
    x = self.fc(x)
    return x

to:

def forward_predictions(self, x):  # x is embedding vector
    x = self.global_pool(x)
    if self.drop_rate:
        x = F.dropout(x, p=float(self.drop_rate), training=self.training)
    x = self.fc(x)
    return x

def forward(self, x):
    x = self.forward_features(x)
    x = self.forward_predictions(x)
    return x

Our inference code, then, would become:

embeddings = model.forward_features(input)
predictions = model.forward_predictions(embeddings)  # no redundant compute
return embeddings, predictions

... and enables us to convert feature embedding vectors from a feature store into a prediction:

embedding = feature_store.query(<interesting image>)
nearest_neighbor_embedding = feature_store.get_nearest_vector(embedding)
prediction = model.forward_predictions(nearest_neighbor_embedding)

Describe alternatives you've considered
Option 1: Two separate forward passes (for embeddings and predictions) for each image. Create "forward_predictions" helper functions for each different architecture we use to convert embeddings into predictions.

Option 2: register a custom forward hook for every architecture which intercepts the embedding vector during the forward pass.

[rwightman]

@crypdick naming semantic debates aside (ie forward_head vs forward_classifier vs forward_predictions) the proposal seems reasonable. But, it would touch every model so I'd like to understand use cases first.

I don't feel it's common to want / need both so curious what uses cases for having both the embeddings / pre-head features and the classifier predictions are? Usually one wants either the predictions or one wants the embeddings to feed through another net / a custom head / or use directly in KNN, distance metrics of some sort, etc

[crypdick]

@rwightman Probably the most common use-case is to plot the embeddings (after reducing dimensionality) and color the points by the predicted class. This is useful to find outliers, or patterns in model mistakes, or systemically mislabeled data, or just to check whether your model has learned a robust internal representation.

Here's a product demo that showcases ways one might use this kind of visualization to improve training data: https://youtu.be/NfaBFDwglUE?t=122

It's common to see this kind of plot presented at the International Conference on Learning Representations and in the representation learning literature generally.

[rwightman]

@crypdick thanks for the info, will check the video soon and think about this a bit ...

I currently have quite a few changes on deck including some that overlap with this (I've shifted the boundary of foward_features for all transformer models to exclude token selection / gap so that they are consistent with CNN, and added ability to change pooling between token and GAP where it makes sense). Essentially not a bad time to consider this....

[crypdick]

For your consideration, I like the idea of every model using a consistent forward pattern, i.e.

def forward(self, x):
    x = self.forward_features(x)
    x = self.forward_head(x)
    return x

In that case, forward_features() would need to pool by default and have an optional unpooled kwarg. My 2c.

Alternatively, something like

def forward(self, x):
    x = self.forward_features(x)
    x = self.forward_pool(x)
    x = self.forward_head(x)
    return x

[rwightman]

@crypdick hmm, the pooling part is tricky... while I can see model.forward_features(unpooled=True) being useful and convenient for some, I feel there are many use cases for features / embedding where the same pooling as the classifier is not desired. It might be better to make that slightly redundant and anyone using forward_features() for embeddings will need to pool themselves.

I often view pooling as being part of the head, there are some heads, even for classification which are based on completely different pooling mechanisms (incl some transformer / self-attn based variants).

[rwightman]

Also, the location of 'pool' can vary and not all models have it nicely between what one would consider 'features' and 'head'

so your idea would need to be closer to below and I gather you want the output of prelogits.

def forward(self, x):
    x = self.forward_features(x)  # always last 2d/1d feature_map/sequence
    x = self.forward_prelogits(x)  # anything between last feature map and last fc layer
    #  (could include other fc, pooling, token selection, norm, etc)
    x = self.forward_head(x)  # final classifier expecting non-spatial input
    return x

[rwightman]

Another possibility that could be done independently of these more significant changes is providing the module names for 'pre logits' so that hooks (or FX extraction) would be more straightforward and common to all models....

[crypdick]

That would work. Consistent naming of these modules (or aliases) would avoid ugly conditional logic, e.g.

        all_embeddings = {}

        def get_activation(name):
            def hook(_model, _input, output):
                all_embeddings[name] = output.detach().cpu()

            return hook
            
        # register hooks
        if architecture.startswith('resnet'):
            hook_handle = self.model.avgpool.register_forward_hook(get_activation("avgpool"))
        elif ...

        with torch.no_grad():
            logits = self.model(imgs).cpu()

        # unregister hook
        hook_handle.remove()

        if architecture.startswith('resnet'):
            embeddings = all_embeddings.get("avgpool", torch.Tensor())
        elif ...

I think that the forward_features, forward_prelogits, forward_head sequence is sensible, and more ergonomic than using PyTorch hooks as above.

[rwightman]

@crypdick K, I can probably do something hook based (see below)

Breaking it down into three (feat, pool/prelogits / head) feels like trouble after more thought, aside from the shifting location of pool, it would likely break some existing usage (verified in downstream code).

forward_features + forward_head would be largely safe, where forward_features would be unpooled, forward_head would do any pooling, extra norm or fv between pool and classifier, and then the classifier. Although this doesn't address your use case directly it would help with it if you pooled yourself, and I'm curious if any others would find it useful? It does actually provide a possibly useful override fn for other use cases that just want to overidde the head forward path...

For your use case, it would be possible to have a helper that leverages existing model info for hooking the input to classifier.

model = EmbeddingsAndPredictions(model)
pre_logits, predictions = model(input)

And if you wanted to do something not model specific right now, you can use forward_pre hooks on the classifier for pre-logits. I believe prehooking model.default_cfg['classifier'] would get you what you want for every model?

[crypdick]

I also like the idea of a EmbeddingsAndPredictions helper function! It's a clean and consistent abstraction.

For clarity's sake, were you thinking something like:

unpooled = model.forward_features(imgs)
gap = model.default_cfg['pooling_layer']
pooled = gap(unpooled)
head = model.default_cfg['classifier']
predictions = head(pooled)

One option, which I think is ugly, is to have the forward_prelogits (or model.default_cfg['pooling_layer']) just be a no-op/identity function in architectures which don't fit the unpooled -> pooled -> predictions mould. That would standardize the interface and enable swapping out models without having to handle edge cases.

[rwightman]

@crypdick pooling isn't consistently a layer across all models, to make that work for the token based ones seemed like more work than just leaving it as code in the fwd

pre-hook (grabbing input to the classifier) is the most consistent way of getting the activations before the classifier regardless of whether it's a module, whether it comes from pooling, etc (no module output to hook)... that's what something like EmbeddingsAndPredictions would use (a pre-hook)

[crypdick]

@rwightman I think you're right-- prehooks are the way to go here.

[rwightman]

@crypdick I've got below implemented in my current set of changes

model = timm.create_model('resnet50')
unpooled_features = model.forward_features(img)
embeddings = model.forward_head(unpool_features, pre_logits=True)
predictions = model.forward_head(unpooled_features)

There is a slight bit of redudancy through the head (usually just pooling twice, but sometimes an extra norm or pre classifier fc/1x1. This will work with scripting and usual eager mode Python. I will also look at a simple hook based wrapper, but it appears there are still challenges getting such a beast to work well with torchscript (sigh).

[crypdick]

Ah! Yeah, I don't know how to make this play nice with torchscript. I think the above solution is ok in most cases. A hook-based solution is nice when the forward_head has non-trivial overhead, for example when using a CNN backbone + a transformer head.