pretrained dino backbone (#303)

* dino implementation
* dino/mae cleanup
* update docs
* fix tests

Author: themattinthehatt

docs/api/lightning_pose.models.backbones.vit_mae.ViTVisionEncoder.rst

@@ -1,16 +1,16 @@
-SamVisionEncoderHF
-==================
+SamVisionEncoder
+================
 
 .. currentmodule:: lightning_pose.models.backbones.vit_sam
 
-.. autoclass:: SamVisionEncoderHF
+.. autoclass:: SamVisionEncoder
    :show-inheritance:
 
    .. rubric:: Methods Summary
 
    .. autosummary::
 
-      ~SamVisionEncoderHF.forward
+      ~SamVisionEncoder.forward
 
    .. rubric:: Methods Documentation
 

…backbones.vit_sam.SamVisionEncoderHF.rst …s.backbones.vit_sam.SamVisionEncoder.rstdocs/api/lightning_pose.models.backbones.vit_sam.SamVisionEncoderHF.rst renamed to docs/api/lightning_pose.models.backbones.vit_sam.SamVisionEncoder.rst docs/api/lightning_pose.models.backbones.vit_sam.SamVisionEncoderHF.rst renamed to docs/api/lightning_pose.models.backbones.vit_sam.SamVisionEncoder.rst

@@ -4,9 +4,6 @@ lightning\_pose.models.backbones
 .. automodapi:: lightning_pose.models.backbones.torchvision
    :no-inheritance-diagram:
 
-.. automodapi:: lightning_pose.models.backbones.vit_mae
-   :no-inheritance-diagram:
-
 .. automodapi:: lightning_pose.models.backbones.vit_sam
    :no-inheritance-diagram:
 

docs/modules/lightning_pose.models.backbones.rst

@@ -207,6 +207,9 @@ The following parameters relate to model architecture and unsupervised losses.
     * efficientnet_b0: EfficientNet-B0 pretrained on ImageNet
     * efficientnet_b1: EfficientNet-B1 pretrained on ImageNet
     * efficientnet_b2: EfficientNet-B2 pretrained on ImageNet
+    * vits_dino: Vision Transformer (Small) pretrained on ImageNet with DINO
+    * vitb_dino: Vision Transformer (Base) pretrained on ImageNet with DINO
+    * vitb_imagenet: Vision Transformer (Base) pretrained on ImageNet with MAE loss
     * vitb_sam: Segment Anything Model (Vision Transformer Base)
 
   Note: the file size for a single ResNet-50 network is approximately 275 MB.

docs/source/user_guide/config_file.rst

@@ -16,6 +16,8 @@
     "efficientnet_b0",
     "efficientnet_b1",
     "efficientnet_b2",
-    "vitb_sam",
+    "vits_dino",
+    "vitb_dino",
     "vitb_imagenet",
+    "vitb_sam",
 ]

lightning_pose/models/backbones/__init__.py

@@ -7,11 +7,11 @@
 
 # to ignore imports for sphix-autoapidoc
 __all__ = [
-    "SamVisionEncoderHF",
+    "SamVisionEncoder",
 ]
 
 
-class SamVisionEncoderHF(nn.Module):
+class SamVisionEncoder(nn.Module):
     """Wrapper around HuggingFace's SAM Vision Encoder."""
 
     def __init__(