SWAG with Vision Transformers

[ayushp88]

Hi, does this repository also contain SWAG with Vision Transformers?

[svirpioj]

Not at the moment. However, it's pretty easy to wrap any model from the transformers library. For example, this should work for ViT:

from transformers import ViTConfig, ViTModel
from swag_transformers.base import SwagConfig, SwagModel

MODEL_TYPE = 'swag_vit'

class SwagViTConfig(SwagConfig):
    """Config for ViT model averaging with SWAG"""

    model_type = MODEL_TYPE
    internal_config_class = ViTConfig

class SwagViTModel(SwagModel):
    """SWAG ViT model"""

    config_class = SwagViTConfig
    base_model_prefix = MODEL_TYPE
    internal_model_class = ViTModel

model = ViTModel.from_pretrained('WinKawaks/vit-tiny-patch16-224')
swag_model = SwagViTModel.from_base(model, no_cov_mat=False)

I can also add new models to the library via pull requests.

[svirpioj]

Now there are basic instructions for this in the README.