TorchEnsemble-Community / Ensemble-Pytorch

A unified ensemble framework for PyTorch to improve the performance and robustness of your deep learning model.
https://ensemble-pytorch.readthedocs.io
BSD 3-Clause "New" or "Revised" License
1.05k stars 95 forks source link

Getting Embeddings from the base model #165

Closed desertSniper87 closed 2 months ago

desertSniper87 commented 2 months ago

Hey everyone!

I've been learning Torchensemble, and I wanted to know how you can extract the embeddings before the final classification layer.

Here is my ensemble model:

VotingClassifier(
 (base_estimator_): CCT(
  (tokenizer): Tokenizer(
   (conv_layers): Sequential(
    (0): Sequential(
     (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
     (1): ReLU()
     (2): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
    )
    (1): Sequential(
     (0): Conv2d(64, 256, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
     (1): ReLU()
     (2): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
    )
   )
   (flattener): Flatten(start_dim=2, end_dim=3)
  )
  (classifier): TransformerClassifier(
   (attention_pool): Linear(in_features=256, out_features=1, bias=True)
   (dropout): Dropout(p=0.0, inplace=False)
   (blocks): ModuleList(
    (0): TransformerEncoderLayer(
     (pre_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
     (self_attn): Attention(
      (qkv): Linear(in_features=256, out_features=768, bias=False)
      (attn_drop): Dropout(p=0.1, inplace=False)
      (proj): Linear(in_features=256, out_features=256, bias=True)
      (proj_drop): Dropout(p=0.0, inplace=False)
     )
     (linear1): Linear(in_features=256, out_features=512, bias=True)
     (dropout1): Dropout(p=0.0, inplace=False)
     (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
     (linear2): Linear(in_features=512, out_features=256, bias=True)
     (dropout2): Dropout(p=0.0, inplace=False)
     (drop_path): Identity()
    )
    (1-6): 6 x TransformerEncoderLayer(
     (pre_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
     (self_attn): Attention(
      (qkv): Linear(in_features=256, out_features=768, bias=False)
      (attn_drop): Dropout(p=0.1, inplace=False)
      (proj): Linear(in_features=256, out_features=256, bias=True)
      (proj_drop): Dropout(p=0.0, inplace=False)
     )
     (linear1): Linear(in_features=256, out_features=512, bias=True)
     (dropout1): Dropout(p=0.0, inplace=False)
     (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
     (linear2): Linear(in_features=512, out_features=256, bias=True)
     (dropout2): Dropout(p=0.0, inplace=False)
     (drop_path): DropPath()
    )
   )
   (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
   (fc): Linear(in_features=256, out_features=5013, bias=True)
  )
 )
 (estimators_): ModuleList()
 (_criterion): CrossEntropyLoss()
)

You see, there is a fc (fully connected layer) in the last of the base model. I want to delete it or replace it with Identity. How to do it? It seems removing the layer in the base estimator does not do anything.

xuyxu commented 2 months ago

Hi, @desertSniper87, estimators are stored in the module estimators_, but it looks like there is nothing in your classifier, could you show us how you initialized the model ?

desertSniper87 commented 2 months ago

Created the ensemble using this

from torchensemble import VotingClassifier

ensemble = VotingClassifier(
    estimator=cct_model,
    n_estimators=5,
    cuda=True
)

Took the cct_model from here:

from src import cct_7_7x2_224
cct_model = cct_14_7x2_224(**model_args)
xuyxu commented 2 months ago

OK, after fitting your ensemble model, you can simply index the target estimator using model = ensemble.estimators_[i], which should be a object of cct_model, you can then extract the embeddings using the corresponding model parameter.

desertSniper87 commented 2 months ago

Thank you for your response.