Cleaned up Vision, added VIT and MAE

dibyaghosh · May 21, 2023 · a5bab4c · a5bab4c
1 parent 6ef3f01
commit a5bab4c
Show file tree

Hide file tree

Showing 8 changed files with 1,301 additions and 1,110 deletions.
diff --git a/jaxrl_m/vision/__init__.py b/jaxrl_m/vision/__init__.py
@@ -1,14 +1,20 @@
 from jaxrl_m.vision.impala import impala_configs
-from jaxrl_m.vision.bigvision_resnetv2 import resnetv2_configs
+from jaxrl_m.vision.resnet_v2 import resnetv2_configs
 from jaxrl_m.vision.small_encoders import small_configs
 from jaxrl_m.vision.bridge_resnet_v1 import bridge_resnetv1_configs
 from jaxrl_m.vision.resnet_v1 import vanilla_resnetv1_configs
 
+from jaxrl_m.vision.mae import mae_model_configs
+from jaxrl_m.vision.vit import vit_configs
+
 from jaxrl_m.vision import data_augmentations
 
 encoders = dict()
 encoders.update(impala_configs)
 encoders.update(resnetv2_configs)
 encoders.update(bridge_resnetv1_configs)
 encoders.update(vanilla_resnetv1_configs)
-encoders.update(small_configs)
+encoders.update(small_configs)
+
+encoders.update(mae_model_configs)
+encoders.update(vit_configs)
diff --git a/jaxrl_m/vision/bigvision_common.py b/jaxrl_m/vision/bigvision_common.py