add training scripts & solve dependency issue

Author: retsuh-bqw

FeedbackPolicy/data/data.py

@@ -1,5 +1,5 @@
+import os
 import torch
-import open_clip
 from transformers import CLIPTextModel, CLIPTokenizer
 
 from visual_planner.trainer import GoalGaussianDiffusion
@@ -11,6 +11,13 @@
 from .vit import VisionTransformer
 
 
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+IMAGENET_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
+IMAGENET_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
+
+
 def load_model(
     clip_vision_encoder_path: str,
     clip_vision_encoder_pretrained: str,
@@ -75,3 +82,52 @@ def load_model(
 
 
     return visual_planner, policy_model, tokenizer, text_encoder
+
+
+
+def create_feedback_policy(
+    vision_encoder: str = 'vc1-base',   #TODO: Support additional visual encoders
+    resume_from_checkpoint: str = None,
+):
+
+    import torchvision.transforms as transforms
+    image_processor = transforms.Compose([
+                            transforms.Resize((192, 192), interpolation = transforms.InterpolationMode.BICUBIC),
+                            transforms.ToTensor(),
+                            transforms.Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD),
+                    ])
+    pretrained_model = "clip-vit-large-patch14"
+    text_tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path = pretrained_model)
+
+    from vc_models.models.vit import model_utils
+    vision_encoder = model_utils.load_model(model_utils.VC1_BASE_NAME)
+    embd_size = 768
+
+
+    model = FeedbackDrivenPolicy(vision_encoder = vision_encoder, \
+            vis_dim = embd_size,
+            window_size = 5,
+            sampling_step = 1)
+    
+    model.vision_encoder.requires_grad_(False)
+
+    def check_file_exists(file_path):
+        if not os.path.isfile(file_path):
+            raise FileNotFoundError(f"The file '{file_path}' does not exist.")
+
+
+    print('Try loading from ckpt')
+    try:
+        check_file_exists(resume_from_checkpoint)
+        old_ckpt = torch.load(resume_from_checkpoint)['model_state_dict']
+
+        # remove 'module.' in original keys
+        new_ckpt = {}
+        for k, v in old_ckpt.items():
+            new_ckpt[k[7:]] = v
+        model.load_state_dict(new_ckpt, strict=False)
+
+    except FileNotFoundError as e:
+        print(e)
+ 
+    return model, image_processor, text_tokenizer

FeedbackPolicy/models/factory.py

@@ -9,8 +9,7 @@
 from einops import rearrange
 from einops import repeat
 
-from .transformer_utils import Block, PatchEmbed, get_2D_position_embeddings,\ 
-                               RMSNorm, SwishGLU
+from .transformer_utils import Block, PatchEmbed, get_2D_position_embeddings, RMSNorm, SwishGLU
 
 
 

FeedbackPolicy/models/policy.py

@@ -0,0 +1,132 @@
+"""
+Util functions for setting up distributed training.
+Credit: https://github.com/mlfoundations/open_clip/blob/main/src/training/distributed.py
+"""
+
+import os
+import torch
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+def is_global_master(args):
+    return args.rank == 0
+
+
+def is_local_master(args):
+    return args.local_rank == 0
+
+
+def is_master(args, local=False):
+    return is_local_master(args) if local else is_global_master(args)
+
+
+def is_using_horovod():
+    # NOTE w/ horovod run, OMPI vars should be set, but w/ SLURM PMI vars will be set
+    # Differentiating between horovod and DDP use via SLURM may not be possible, so horovod arg still required...
+    ompi_vars = ["OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"]
+    pmi_vars = ["PMI_RANK", "PMI_SIZE"]
+    if all([var in os.environ for var in ompi_vars]) or all(
+        [var in os.environ for var in pmi_vars]
+    ):
+        return True
+    else:
+        return False
+
+
+def is_using_distributed():
+    if "WORLD_SIZE" in os.environ:
+        return int(os.environ["WORLD_SIZE"]) > 1
+    if "SLURM_NTASKS" in os.environ:
+        return int(os.environ["SLURM_NTASKS"]) > 1
+    return False
+
+
+def world_info_from_env():
+    local_rank = 0
+    for v in (
+        "LOCAL_RANK",
+        "MPI_LOCALRANKID",
+        "SLURM_LOCALID",
+        "OMPI_COMM_WORLD_LOCAL_RANK",
+    ):
+        if v in os.environ:
+            local_rank = int(os.environ[v])
+            break
+    global_rank = 0
+    for v in ("RANK", "PMI_RANK", "SLURM_PROCID", "OMPI_COMM_WORLD_RANK"):
+        if v in os.environ:
+            global_rank = int(os.environ[v])
+            break
+    world_size = 1
+    for v in ("WORLD_SIZE", "PMI_SIZE", "SLURM_NTASKS", "OMPI_COMM_WORLD_SIZE"):
+        if v in os.environ:
+            world_size = int(os.environ[v])
+            break
+
+    return local_rank, global_rank, world_size
+
+
+def init_distributed_device(args):
+    # Distributed training = training on more than one GPU.
+    # Works in both single and multi-node scenarios.
+    args.distributed = False
+    args.world_size = 1
+    args.rank = 0  # global rank
+    args.local_rank = 0
+    if args.horovod:
+        assert hvd is not None, "Horovod is not installed"
+        hvd.init()
+        args.local_rank = int(hvd.local_rank())
+        args.rank = hvd.rank()
+        args.world_size = hvd.size()
+        args.distributed = True
+        os.environ["LOCAL_RANK"] = str(args.local_rank)
+        os.environ["RANK"] = str(args.rank)
+        os.environ["WORLD_SIZE"] = str(args.world_size)
+    elif is_using_distributed():
+        if "SLURM_PROCID" in os.environ:
+            # DDP via SLURM
+            args.local_rank, args.rank, args.world_size = world_info_from_env()
+            # SLURM var -> torch.distributed vars in case needed
+            os.environ["LOCAL_RANK"] = str(args.local_rank)
+            os.environ["RANK"] = str(args.rank)
+            os.environ["WORLD_SIZE"] = str(args.world_size)
+            torch.distributed.init_process_group(
+                backend=args.dist_backend,
+                init_method=args.dist_url,
+                world_size=args.world_size,
+                rank=args.rank,
+            )
+        else:
+            # DDP via torchrun, torch.distributed.launch
+            args.local_rank, _, _ = world_info_from_env()
+            torch.distributed.init_process_group(
+                backend=args.dist_backend, init_method=args.dist_url
+            )
+            args.world_size = torch.distributed.get_world_size()
+            args.rank = torch.distributed.get_rank()
+        args.distributed = True
+    else:
+        # needed to run on single gpu
+        torch.distributed.init_process_group(
+            backend=args.dist_backend,
+            init_method=args.dist_url,
+            world_size=1,
+            rank=0,
+        )
+
+    if torch.cuda.is_available():
+        if args.distributed and not args.no_set_device_rank:
+            device = "cuda:%d" % args.local_rank
+        else:
+            device = "cuda:0"
+        torch.cuda.set_device(device)
+    else:
+        device = "cpu"
+    args.device = device
+    device = torch.device(device)
+    return device