Merge pull request octo-models#108 from rail-berkeley/kevin-tmp-shuffle

Shuffle before decode, interleave at transition level

Author: kvablack

config.py

@@ -66,15 +66,15 @@ def get_config(
                 clip_gradient=1.0,
                 frozen_keys=tuple(),
             ),
-            batch_size=1024,
             eval_batch_size=128,
-            shuffle_buffer_size=100000,
+            prefetch_num_batches=0,
             val_shuffle_buffer_size=1000,
             num_val_batches=16,
             start_step=placeholder(int),
             log_interval=100,
             eval_interval=5000,
-            save_interval=5000,
+            viz_interval=20000,
+            save_interval=10000,
             trajs_for_metrics=100,
             trajs_for_viz=8,
             resume_path=placeholder(str),
@@ -114,7 +114,7 @@ def get_dataset_config(modality="multimodal", window_size=1):
         raise ValueError(f"Unknown modality {modality}")
 
     return {
-        # oxe_kwargs will generate data_kwargs_list and sampling weights
+        # oxe_kwargs will generate dataset_kwargs_list and sampling weights
         "oxe_kwargs": dict(
             data_mix=placeholder(str),
             # for v4 TPUs: "gs://rail-orca-central2/resize_336_336"
@@ -123,18 +123,19 @@ def get_dataset_config(modality="multimodal", window_size=1):
             n_wrist_cameras=0,
             load_depth=False,
         ),
-        # common_kwargs override specific kwargs from data_kwargs_list
-        "common_kwargs": dict(
-            ram_budget=1,  # limit RAM per dataset
-            num_parallel_reads=8,  # for reading from GCS
-            num_parallel_calls=16,  # for the less CPU-intensive ops in initial dataset construction
+        # common_dataset_kwargs override specific kwargs from dataset_kwargs_list
+        "common_dataset_kwargs": dict(
             action_proprio_normalization_type=normalization_type,
         ),
-        "transform_kwargs": dict(
-            resize_size=(256, 256),
-            num_parallel_calls=32,  # for the most CPU-intensive ops (decoding, resizing, augmenting)
+        "traj_transform_kwargs": dict(
             window_size=window_size,
             additional_action_window_size=0,
+            goal_relabeling_strategy="uniform",
+            subsample_length=100,
+            **task_augmentation,
+        ),
+        "frame_transform_kwargs": dict(
+            resize_size=(256, 256),
             image_augment_kwargs=dict(
                 random_resized_crop=dict(scale=[0.8, 1.0], ratio=[0.9, 1.1]),
                 random_brightness=[0.2],
@@ -149,9 +150,13 @@ def get_dataset_config(modality="multimodal", window_size=1):
                     "random_hue",
                 ],
             ),
-            goal_relabeling_strategy="uniform",
-            **task_augmentation,
         ),
+        "traj_transform_threads": 48,  # shared between all datasets
+        "traj_read_threads": 48,  # shared between all datasets
+        "frame_transform_threads": 200,  # not shared between datasets
+        "shuffle_buffer_size": 100000,  # shared between all datasets
+        "batch_size": 1024,
+        "balance_weights": True,
     }
 
 

experiments/dibya/finetune_config.py

@@ -32,7 +32,6 @@ def get_config(
         "action_encoding": ActionEncoding.EEF_POS,
         "action_proprio_normalization_type": "normal",
         # If the default data loading speed is too slow, try these:
-        # and "num_parallel_calls" in `transform_kwargs` below
         # "num_parallel_reads": 8,  # for reading from disk / GCS
         # "num_parallel_calls": 16,  # for initial dataset construction
     }
@@ -70,7 +69,7 @@ def get_config(
         wandb=dict(
             project="orca_finetune", group=placeholder(str), entity=placeholder(str)
         ),
-        finetuning_dataset=FINETUNING_KWARGS,
+        dataset_kwargs=FINETUNING_KWARGS,
         modality=task,
         finetuning_mode=mode,
         window_size=window_size,
@@ -107,9 +106,18 @@ def get_config(
     else:
         raise ValueError("Invalid modality")
 
-    transform_kwargs = dict(
+    traj_transform_kwargs = dict(
         window_size=window_size,
         additional_action_window_size=0,
+        goal_relabeling_strategy=goal_relabeling_strategy,
+        task_augmentation_strategy="delete_task_conditioning",
+        task_augmentation_kwargs=dict(
+            delete_key_groups_probs=delete_key_groups_probs,
+        ),
+        # If the default data loading speed is too slow, try these:
+        # num_parallel_calls=16,  # for less CPU-intensive ops
+    )
+    frame_transform_kwargs = dict(
         resize_size=(256, 256),
         image_augment_kwargs=dict(
             random_resized_crop=dict(scale=[0.8, 1.0], ratio=[0.9, 1.1]),
@@ -125,13 +133,12 @@ def get_config(
                 "random_hue",
             ],
         ),
-        goal_relabeling_strategy=goal_relabeling_strategy,
-        task_augmentation_strategy="delete_task_conditioning",
-        task_augmentation_kwargs=dict(
-            delete_key_groups_probs=delete_key_groups_probs,
-        ),
-        # If the default data loading speed is too slow, try these:
-        # num_parallel_calls=16,  # for the most CPU-intensive ops (decoding, resizing, augmenting)
     )
-    config["data_transforms"] = transform_kwargs
+    # If the default data loading speed is too slow, try these:
+    config[
+        "frame_transform_threads"
+    ] = 16  # for the most CPU-intensive ops (decoding, resizing, augmenting)
+
+    config["traj_transform_kwargs"] = traj_transform_kwargs
+    config["frame_transform_kwargs"] = frame_transform_kwargs
     return ConfigDict(config)

finetune.py

@@ -64,8 +64,8 @@ def main(_):
         ORCA Finetuning Script
         ======================
         Pretrained model: {FLAGS.config.pretrained_path}
-        Finetuning Dataset: {FLAGS.config.finetuning_dataset.name}
-        Data dir: {FLAGS.config.finetuning_dataset.data_dir}
+        Finetuning Dataset: {FLAGS.config.dataset_kwargs.name}
+        Data dir: {FLAGS.config.dataset_kwargs.data_dir}
         Task Modality: {FLAGS.config.modality}
         Finetuning Mode: {FLAGS.config.finetuning_mode}
 
@@ -159,13 +159,22 @@ def process_text(batch):
         del batch["dataset_name"]
         return batch
 
-    data_kwargs = FLAGS.config.finetuning_dataset
-    transform_kwargs = FLAGS.config.data_transforms
-
-    dataset = make_single_dataset(data_kwargs, transform_kwargs, train=True)
-    val_dataset = make_single_dataset(data_kwargs, transform_kwargs, train=False)
+    dataset = make_single_dataset(
+        FLAGS.config.dataset_kwargs,
+        FLAGS.config.traj_transform_kwargs,
+        FLAGS.config.frame_transform_kwargs,
+        train=True,
+        frame_transform_threads=FLAGS.config.frame_transform_threads,
+    )
+    val_dataset = make_single_dataset(
+        FLAGS.config.dataset_kwargs,
+        FLAGS.config.traj_transform_kwargs,
+        FLAGS.config.frame_transform_kwargs,
+        train=False,
+        frame_transform_threads=FLAGS.config.frame_transform_threads,
+    )
     visualizer = Visualizer(
-        val_dataset, text_processor=text_processor, cache_trajs=False
+        val_dataset, text_processor=text_processor, freeze_trajs=False
     )
 
     def create_iterator(dataset):