ViT Image Classification Support

[tanmoyio]

Starter code

import pytorch_lightning as pl
from transformers import AutoFeatureExtractor

from lightning_transformers.task.vision.image_classification import (
    ImageClassificationDataConfig,
    ImageClassificationDataModule,
    ImageClassificationTransformer,
)

MODEL_NAME = "nateraw/vit-base-beans"
feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path=MODEL_NAME)

dm = ImageClassificationDataModule(
    cfg=ImageClassificationDataConfig(
        batch_size=8,
        dataset_name="beans",
        num_workers=8
    ),
    feature_extractor=feature_extractor,
)

model = ImageClassificationTransformer(pretrained_model_name_or_path=MODEL_NAME)
trainer = pl.Trainer(accelerator="gpu", max_epochs=5)
trainer.fit(model, dm)

[codecov[bot]]

Codecov Report

Merging #252 (b7d315b) into master (4832f00) will increase coverage by 1%. The diff coverage is 93%.

@@          Coverage Diff          @@
##           master   #252   +/-   ##
=====================================
+ Coverage      74%    75%   +1%     
=====================================
  Files          69     73    +4     
  Lines        1534   1622   +88     
=====================================
+ Hits         1128   1210   +82     
- Misses        406    412    +6     

[SeanNaren]

Added some simple tests, the example and we're good to go I think! this is epic, thanks so much @tanmoyio for your contribution ❤️

I'm going to assume the VIT model is something you'd like to use, but if you're interested in tackling a Wav2vec task as well it would be great to get your assistance on: https://github.com/PyTorchLightning/lightning-transformers/pull/251

[SeanNaren]

I'll make a separate issue to add some docs in a separate PR.

[tanmoyio]

@SeanNaren Yeah I would love to add wav2vec2 tasks, may be in another category speech

[SeanNaren]

@SeanNaren Yeah I would love to add wav2vec2 tasks, may be in another category speech

that sounds perfect! agreed, I think the PR that was made requires a lot of work, might be worth making it from scratch.

[SeanNaren]

All fixed! merging this down, thanks so much again 🚀