Add code for training

.gitignore

@@ -1,5 +1,6 @@
 # Additional 
 configs/user/nguyen.yaml
+sync_wandb.sh
 .hydra/
 outputs/
 None/

README.md

@@ -130,24 +130,29 @@ If everything is fine, here are the number of images that you should get:
  ##  Launch a training  :rocket:
 
 <details><summary>Click to expand</summary>
+
 ### 0. (Optional) We use pretrained weight from MoCo v2. You can download it from [here]() or run:
+
 ```
 python -m src.scripts.download_moco_weights
 ```
+
 If you don't want to use pretrained weights, you can remove the path in [this line](https://drive.google.com/drive/folders/1p9eJ8dTxR3rVinvaFxPw5N_3IGSlS2_E?usp=sharing).
 ### 1. Training on all BOP datasets except LINEMOD and T-LESS (only objects 19-30)
 ```
 python train.py name_exp=train_all
 ```
 
 The parsing is done with Hydra library. You can override anything in the configuration by passing arguments. For example:
+
 ```
 # experiment 1: change batch_size, using data augmentation, update name_exp
 python train.py machine.batch_size=2 use_augmentation=True name_exp=train_augmentation
 
 # experiment 2: change batch_size, using data augmentation, update name_exp, update_lr
 python train.py machine.batch_size=2 use_augmentation=True model.lr=0.001 name_exp=train_augmentation_lr0.001
 ```
+
 </details>
 
 ## Acknowledgement

configs/data/hb.yaml

@@ -11,7 +11,7 @@ dataloader:
   _target_: src.dataloader.bop.BOPDataset
   root_dir: ${machine.root_dir}/datasets/hb
   template_dir: ${machine.root_dir}/datasets/templates/hb
-  obj_ids: 
+  obj_ids: "1,3,4,5,8,10,11,12,13,14,15,16,17,18,19,20,22,23,24,25,26,27,28,29,30,31,32,33"
   reset_metaData: True
   img_size: 256
   cropping_with_bbox: True

configs/data/tless_train.yaml

@@ -11,7 +11,7 @@ dataloader:
   _target_: src.dataloader.bop.BOPDataset
   root_dir: ${machine.root_dir}/datasets/tless/train
   template_dir: ${machine.root_dir}/datasets/templates/tless
-  obj_ids:
+  obj_ids: "1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18"
   reset_metaData: True
   img_size: 256
   cropping_with_bbox: True

configs/model/resnet50.yaml

@@ -11,4 +11,4 @@ weight_decay: 0.0005
 warm_up_steps: 1000
 log_interval: 1000
 log_dir: ${save_dir}/media
-use_all_gather: false
+use_all_gather: true # for multi-gpu training with larger batch_size -> efficient for contrast learning with InfoNCE

src/dataloader/base.py

@@ -56,13 +56,13 @@ def load_list_scene(self, split=None):
                 ]
             )
         elif isinstance(split, list):
-            self.list_scenes = sorted(
-                [
-                    osp.join(self.root_dir, scene)
-                    for scene in split
-                    if os.path.isdir(osp.join(self.root_dir, scene))
-                ]
-            )
+            self.list_scenes = []
+            for scene in split:
+                if not isinstance(scene, str):
+                    scene = f"{scene:06d}"
+                if os.path.isdir(osp.join(self.root_dir, scene)):
+                    self.list_scenes.append(osp.join(self.root_dir, scene))
+            self.list_scenes = sorted(self.list_scenes)
         else:
             raise NotImplementedError
         logging.info(f"Found {len(self.list_scenes)} scenes")

src/dataloader/bop.py

@@ -41,7 +41,10 @@ def __init__(
         self.augmentator = Augmentator()
 
         self.load_template_poses(template_dir=template_dir)
-        self.load_list_scene(split=obj_ids if obj_ids is not None else split)
+        if isinstance(obj_ids, str):
+            obj_ids = [int(obj_id) for obj_id in obj_ids.split(",")]
+            logging.info(f"ATTENTION: Loading {len(obj_ids)} objects!")
+        self.load_list_scene(split=split)
         self.load_metaData(
             reset_metaData=reset_metaData,
             mode="query",
@@ -74,7 +77,9 @@ def __init__(
                 transforms.Lambda(lambda mask: torch.from_numpy(mask).unsqueeze(0)),
             ]
         )
-        logging.info(f"Length of dataloader: {self.__len__()}")
+        logging.info(
+            f"Length of dataloader: {self.__len__()} containing objects {np.unique(self.metaData['obj_id'])}"
+        )
 
     def load_template_poses(self, template_dir):
         self.templates_poses = np.load(osp.join(template_dir, "obj_poses.npy"))
@@ -85,18 +90,21 @@ def subsample(self, df, percentage):
         selected_obj_id = [id for id in self.obj_ids]
         logging.info(f"Available {avail_obj_id}, selected {selected_obj_id} ")
         selected_index = []
-        for obj_id in self.obj_ids:
-            df_obj = df[df["obj_id"] == obj_id]
-            df_obj = df[df["visib_fract"] == 1]
+        index_dataframe = np.arange(0, len(df))
+        for obj_id in selected_obj_id:
+            selected_index_obj = index_dataframe[# df["obj_id"] == obj_id]
+                np.logical_and(df["obj_id"] == obj_id, df["visib_fract"] >= 0.5)]
             if percentage > 50:
-                df_obj = df_obj[: int(percentage / 100 * len(df_obj))]  # keep first
+                selected_index_obj = selected_index_obj[
+                    : int(percentage / 100 * len(selected_index_obj))
+                ]  # keep first
             else:
-                df_obj = df_obj[
-                    int((1 - percentage / 100) * len(df_obj)) :
+                selected_index_obj = selected_index_obj[
+                    int((1 - percentage / 100) * len(selected_index_obj)) :
                 ]  # keep last
-            selected_index.extend(df_obj.index)
-        df = df.loc[selected_index]
-        logging.info(f"Subsampled to {len(df)} ({percentage}%) images")
+            selected_index.extend(selected_index_obj.tolist())
+        df = df.iloc[selected_index]
+        logging.info(f"Subsampled from {len(index_dataframe)} to {len(df)} ({percentage}%) images")
         return df
 
     def __len__(self):

src/scripts/sanity_check_datasets.py

@@ -60,8 +60,8 @@ def render(cfg: DictConfig) -> None:
             )
             cad_paths.append(cad_path)
             template_dirs.append(template_dir)
-
     os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpus
+
     start_time = time.time()
     pool = multiprocessing.Pool(processes=int(cfg.num_workers))
     call_blender_proc_with_index = partial(
@@ -81,7 +81,7 @@ def render(cfg: DictConfig) -> None:
         )
     )
     finish_time = time.time()
-    print("Total time to render templates for query:", finish_time - start_time)
+    print(f"Total time to render templates for query: {finish_time - start_time}")
 
 
 if __name__ == "__main__":

train.py

@@ -82,6 +82,7 @@ def train(cfg: DictConfig):
         logging.info(
             f"Loading train dataloader with {data_name}, size {len(train_dataloader)} done!"
         )
+        logging.info("---"*100)
         train_dataloaders[data_name] = train_dataloader
     from src.utils.dataloader import concat_dataloader