[HYP] add save_hub.

.gitignore

@@ -16,7 +16,6 @@ notebooks
 .hydra
 old
 .idea
-hub
 *old
 lightning_logs
 results_neurips

bin/banana/banana_RD.sh

@@ -37,7 +37,7 @@ general_kwargs="
 is_only_feat=False
 featurizer=neural_feat
 optimizer@optimizer_feat=Adam
-optimizer_feat.kwargs.lr=1e-3
+optimizer_feat.kwargs.lr=3e-4
 scheduler@scheduler_feat=expdecay1000
 optimizer@optimizer_coder=Adam
 scheduler@scheduler_coder=expdecay100
@@ -47,7 +47,7 @@ trainer.precision=32
 architecture@predictor=mlp_probe
 optimizer@optimizer_pred=Adam
 scheduler@scheduler_pred=unifmultistep100
-optimizer_pred.kwargs.lr=1e-3
+optimizer_pred.kwargs.lr=3e-4
 featurizer.loss.beta_anneal=constant
 "
 

bin/banana/banana_viz_VIC.sh

@@ -38,7 +38,7 @@ general_kwargs="
 is_only_feat=False
 featurizer=neural_feat
 optimizer@optimizer_feat=Adam
-optimizer_feat.kwargs.lr=1e-3
+optimizer_feat.kwargs.lr=3e-4
 scheduler@scheduler_feat=expdecay1000
 optimizer@optimizer_coder=Adam
 scheduler@scheduler_coder=expdecay100
@@ -48,8 +48,9 @@ trainer.precision=32
 architecture@predictor=mlp_probe
 optimizer@optimizer_pred=Adam
 scheduler@scheduler_pred=unifmultistep100
-optimizer_pred.kwargs.lr=1e-3
+optimizer_pred.kwargs.lr=3e-4
 featurizer.loss.beta_anneal=constant
+encoder.z_dim=2
 "
 
 kwargs="
@@ -66,10 +67,14 @@ $add_kwargs
 kwargs_multi="
 data@data_feat=banana_rot
 featurizer.loss.beta=0.07
+distortion=VIC
+encoder.z_dim=1,2
 "
+#VAE
+#encoder.z_dim=1,2
 
 if [ "$is_plot_only" = false ] ; then
-  for kwargs_dep in  "distortion=VAE encoder.z_dim=2"  "distortion=VIC encoder.z_dim=1" 
+  for kwargs_dep in  "" 
   do
 
     python "$main" +hydra.job.env_set.WANDB_NOTES="\"${notes}\"" $kwargs $kwargs_multi $kwargs_dep -m &

bin/banana/banana_viz_VIC_trnslt.sh

@@ -37,7 +37,7 @@ general_kwargs="
 is_only_feat=False
 featurizer=neural_feat
 optimizer@optimizer_feat=Adam
-optimizer_feat.kwargs.lr=1e-3
+optimizer_feat.kwargs.lr=3e-4
 scheduler@scheduler_feat=expdecay1000
 optimizer@optimizer_coder=Adam
 scheduler@scheduler_coder=expdecay100
@@ -47,7 +47,7 @@ trainer.precision=32
 architecture@predictor=mlp_probe
 optimizer@optimizer_pred=Adam
 scheduler@scheduler_pred=unifmultistep100
-optimizer_pred.kwargs.lr=1e-3
+optimizer_pred.kwargs.lr=3e-4
 featurizer.loss.beta_anneal=constant
 "
 

bin/clip/clip_bottleneck_linear_eval.sh

@@ -32,7 +32,7 @@ data@data_pred=stl10,cars196,caltech101,food101,pcam,pets37,cifar10,cifar100,ima
 " 
 
 if [ "$is_plot_only" = false ] ; then
-  for beta in    "1e-1"  "5e-2"  "1e-2"
+  for beta in  "1e-01"   "5e-02"    "1e-02"    
   do
 
     python utils/Z_linear_eval.py  $kwargs $kwargs_multi featurizer.loss.beta=$beta paths.pretrained.load=$pretrained_path/beta$beta  $kwargs_dep -m &

bin/clip/clip_bottleneck_mlp_eval.sh

@@ -60,7 +60,7 @@ seed=int(interval(0,10))
 if [ "$is_plot_only" = false ] ; then
   for data in "stl10" "imagenet" "cars196"  "caltech101"  "food101"  "pcam" "pets37" "cifar10"    "cifar100"        
   do
-    for beta in     "1e-1"  "5e-2"   "1e-2"        
+    for beta in  "1e-01"   "5e-02"    "1e-02"           
     do
 
     python "$main" +hydra.job.env_set.WANDB_NOTES="\"${notes}\"" $kwargs $kwargs_multi data@data_pred=$data featurizer.loss.beta=$beta paths.pretrained.load=$pretrained_path/beta$beta hydra.sweeper.study_name=$data_$beta -m &

bin/clip/clip_bottleneck_pretrain.sh

@@ -29,7 +29,7 @@ $add_kwargs
 kwargs_multi="" 
 
 if [ "$is_plot_only" = false ] ; then
-  for beta in  "1e-1"   "5e-2"    "1e-2"          
+  for beta in  "1e-01"   "5e-02"    "1e-02"          
   do
 
     python "$main" +hydra.job.env_set.WANDB_NOTES="\"${notes}\"" $kwargs $kwargs_multi featurizer.loss.beta=$beta paths.pretrained.save=$pretrained_path/beta$beta -m &

bin/clip/clip_hub.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+
+experiment="clip_hub"
+notes="
+**Goal**: Save all pretrained models to pytorch hub.
+"
+
+# parses special mode for running the script
+source `dirname $0`/../utils.sh
+
+SCRIPT=`realpath $0`
+SCRIPTPATH=`dirname $SCRIPT`
+pretrained_path="$SCRIPTPATH"/../../hub
+
+
+# define all the arguments modified or added to `conf`. If they are added use `+`
+kwargs="
+experiment=$experiment 
+timeout=$time
+encoder.z_dim=512
+data@data_feat=coco
+featurizer=bottleneck_clip_lossyZ
+checkpoint@checkpoint_pred=bestValLoss
+featurizer.is_train=false
+evaluation.communication.ckpt_path=null
+$add_kwargs
+"
+
+for beta in  "1e-01"   "5e-02"    "1e-02"          
+  do
+
+    col_val_subset=""
+       python utils/save_hub.py \
+       load_pretrained.experiment=$experiment  \
+       $col_val_subset \
+       $kwargs  \
+       featurizer.loss.beta=$beta \
+       paths.pretrained.load=$pretrained_path/beta$beta \
+       server=local \
+       trainer.gpus=0 \
+       load_pretrained.mode=[] 
+
+  done
+

bin/mnist/augmnist_viz_VIC.sh

@@ -27,11 +27,6 @@ $add_kwargs
 kwargs_multi="
 distortion=VIC,VAE
 " 
-kwargs_multi="
-distortion=VIC
-optimizer_feat.kwargs.lr=1e-3,1e-4,3e-4
-trainer.max_epochs=20
-" 
 
 if [ "$is_plot_only" = false ] ; then
   for kwargs_dep in  "" 

config/main.yaml

@@ -14,7 +14,7 @@ defaults:
   - architecture@encoder: resnet18
   - architecture@online_evaluator: mlp_probe
   - checkpoint@checkpoint_feat: bestValLoss # NB : will not be able to use Val loss for INCE in the case where you don't use data augmentations at test time because validation is meaningless for distortion. 
-  - optimizer@optimizer_feat: AdamW_lr3e-4_w1e-5 # larger learning rate gives smaller rate but smaller acc. should stick to using adam if not changing the  beta will effectively change the lr
+  - optimizer@optimizer_feat: AdamW_lr1e-3_w1e-5 # larger learning rate gives smaller rate but smaller acc. should stick to using adam if not changing the  beta will effectively change the lr
   - scheduler@scheduler_feat: expdecay100 
   - optimizer@optimizer_coder: AdamW_lr3e-4_w1e-5 # quite robust to this but lr shouldn't be to small
   - scheduler@scheduler_coder: expdecay100

config/optimizer/AdamW_lr1e-3_w1e-5.yaml

@@ -0,0 +1,5 @@
+# name is automatic
+mode: AdamW
+kwargs:
+  lr: 1e-3
+  weight_decay: 1e-5

utils/aggregate.py

@@ -40,6 +40,7 @@
     cfg_save,
     getattr_from_oneof,
     omegaconf2namespace,
+    format_resolver,
 )
 from utils.postplotting import (  # isort:skip
     PRETTY_RENAMER,
@@ -955,4 +956,5 @@ def is_pareto_optimal(costs):
 
 
 if __name__ == "__main__":
+    OmegaConf.register_new_resolver("format", format_resolver)
     main_cli()

utils/save_hub.py

@@ -0,0 +1,57 @@
+"""Save the rate estimator from a pretrained model to hub.
+
+This should be called by `python utils/save_hub.py <conf>` where <conf> sets all configs used to 
+pretrain the model in the first place`.
+"""
+import logging
+import os
+import sys
+from pathlib import Path
+
+import torch
+
+import hydra
+from omegaconf import OmegaConf
+
+MAIN_DIR = os.path.abspath(str(Path(__file__).parents[1]))
+CURR_DIR = os.path.abspath(str(Path(__file__).parents[0]))
+sys.path.append(MAIN_DIR)
+sys.path.append(CURR_DIR)
+
+
+from utils.load_pretrained import PretrainedAnalyser  # isort:skip
+from utils.helpers import format_resolver  # isort:skip
+
+logger = logging.getLogger(__name__)
+
+
+@hydra.main(config_path=f"{MAIN_DIR}/config", config_name="load_pretrained")
+def main_cli(cfg):
+    # uses main_cli sot that `main` can be called from notebooks.
+    try:
+        return main(cfg)
+    except:
+        logger.exception("Failed to save pretrained with this error:")
+        # don't raise error because if multirun want the rest to work
+        pass
+
+
+def main(cfg):
+
+    analyser = PretrainedAnalyser(**cfg.load_pretrained.kwargs)
+
+    logger.info(f"Collecting the data ..")
+    analyser.collect_data(cfg, **cfg.load_pretrained.collect_data)
+
+    rate_estimator = analyser.modules["featurizer"].rate_estimator
+    model_name = f"{MAIN_DIR}/hub/beta{cfg.featurizer.loss.beta:.0e}/factorized_rate.pt"
+    torch.save(
+        rate_estimator.state_dict(), model_name,
+    )
+
+    logger.info(f"Saved pretrained model to {model_name}...")
+
+
+if __name__ == "__main__":
+    OmegaConf.register_new_resolver("format", format_resolver)
+    main_cli()