minor debug tweaks

Author: Erotemic

train_kwcoco_demo.sh

@@ -85,6 +85,7 @@ TRAIN_DPATH="$BUNDLE_DPATH/kwcoco-demo-train-dir"
 # This might only work in development mode, otherwise we will get site packages
 # That still might be fine, but we do want to fix this to run anywhere.
 cd "$REPO_DPATH"
+export CUDA_VISIBLE_DEVICES="1,"
 LOG_BATCH_VIZ_TO_DISK=1 python -m yolo.lazy \
     task=train \
     dataset=kwcoco-demo \

yolo/model/yolo.py

@@ -233,11 +233,13 @@ def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True,
     Returns:
         YOLO: An instance of the model defined by the given configuration.
     """
+    logger.info = print
+    logger.info('CREATE MODEL')
     OmegaConf.set_struct(model_cfg, False)
     model = YOLO(model_cfg, class_num)
     if weight_path:
         logger.info('🏋 Initializing weights')
-        if weight_path == True:
+        if weight_path is True:
             weight_path = Path("weights") / f"{model_cfg.name}.pt"
         elif isinstance(weight_path, str):
             weight_path = Path(weight_path)
@@ -251,4 +253,5 @@ def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True,
             logger.info(":white_check_mark: Success load model & weight")
     else:
         logger.info(":white_check_mark: Success load model without weights")
+    logger.info('CREATED MODEL')
     return model

yolo/tools/data_loader.py

@@ -51,7 +51,7 @@ def load_data(self, dataset_path: Path, phase_name: str):
         Returns:
             dict: The loaded data from the cache for the specified phase.
         """
-        cache_path = dataset_path / f"{phase_name}.cache"
+        cache_path = dataset_path / f"{phase_name}-v1.cache"
 
         if not cache_path.exists():
             logger.info(f":factory: Generating {phase_name} cache")
@@ -232,6 +232,7 @@ def get_data(self, idx):
         img_path, bboxes = self.img_paths[idx], self.bboxes[idx]
         valid_mask = bboxes[:, 0] != -1
 
+        # TODO: we can load an overview here to make this much more efficent
         USE_OVERVIEW_HACK = 0
         if USE_OVERVIEW_HACK:
             # Can leverage overviews to load images faster if they exist.
@@ -252,8 +253,9 @@ def get_data(self, idx):
             with Image.open(img_path) as img:
                 img = img.convert("RGB")
 
-        # TODO: we can load an overview here to make this much more efficent
-        return img, torch.from_numpy(bboxes[valid_mask]), img_path
+        valid_boxes = bboxes[valid_mask]
+        valid_boxes = torch.from_numpy(valid_boxes)
+        return img, valid_boxes, img_path
 
     def get_more_data(self, num: int = 1):
         indices = torch.randint(0, len(self), (num,))

yolo/utils/logging_utils.py

@@ -253,7 +253,9 @@ def on_validation_batch_end(self, trainer: Trainer, pl_module, outputs, batch, b
                 true_dets = tensor_to_kwimage(gt_boxes).numpy()
                 pred_dets = tensor_to_kwimage(pred_boxes).numpy()
                 pred_dets = pred_dets.non_max_supress(thresh=0.3)
-                # pred_dets = pred_dets.compress(pred_dets.scores > 0.1)
+                pred_dets_2 = pred_dets.compress(pred_dets.scores > 0.01)
+                if len(pred_dets_2) > 0:
+                    pred_dets = pred_dets_2
 
                 raw_canvas = image_hwc.copy()
                 true_canvas = true_dets.draw_on(raw_canvas.copy(), color='green')
@@ -266,7 +268,7 @@ def on_validation_batch_end(self, trainer: Trainer, pl_module, outputs, batch, b
                     raw_canvas, true_canvas, pred_canvas
                 ], axis=1, pad=3)
 
-                fname = f'img_{epoch:04d}_{batch_idx:04d}.jpg'
+                fname = f'img_{epoch:04d}_{bx:04d}.jpg'
                 fpath = out_dpath / fname
                 kwimage.imwrite(fpath, canvas)
 
@@ -357,7 +359,9 @@ def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=Tru
 
     print(f'cfg.use_tensorboard={cfg.use_tensorboard}')
     if cfg.use_tensorboard:
-        loggers.append(TensorBoardLogger(log_graph="all", save_dir=save_path))
+        print(f'save_path={save_path}')
+        # loggers.append(TensorBoardLogger(log_graph="all", save_dir=save_path))
+        loggers.append(TensorBoardLogger(save_path))
         from yolo.utils.callbacks.tensorboard_plotter import TensorboardPlotter
         callbacks.append(TensorboardPlotter())
     if cfg.use_wandb: