added working validation

run_val.py

@@ -1,72 +1,32 @@
-from ultralytics import YOLOv10, YOLO
+from ultralytics import YOLOv10, YOLO, YOLOv10PGT
 # from ultralytics.engine.pgt_trainer import PGTTrainer
-# from ultralytics import BaseTrainer
-# from ultralytics.engine.trainer import BaseTrainer
 import os
+from ultralytics.models.yolo.segment import PGTSegmentationTrainer
+import argparse
+from datetime import datetime
 
-# Set CUDA device (only needed for multi-gpu machines) 
+# nohup python run_pgt_train.py --device 1 > ./output_logs/gpu1_yolov10_pgt_train.log 2>&1 & 
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 
-os.environ["CUDA_VISIBLE_DEVICES"] = "4" 
 
-# model = YOLOv10()
+def main(args):
-# model = YOLO()
-# If you want to finetune the model with pretrained weights, you could load the 
-# pretrained weights like below
-# model = YOLOv10.from_pretrained('jameslahm/yolov10{n/s/m/b/l/x}')
-# or
-# wget https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10{n/s/m/b/l/x}.pt
-model = YOLOv10('yolov10n.pt')
 
-# Evaluate the model on the validation set
+    model = YOLOv10PGT(args.model_path)
-results = model.val(data='coco.yaml')
     
-# Print the evaluation results
+    # Evaluate the model on the validation set
-print(results)
+    results = model.val(data=args.data_yaml)
     
-# pred = results[0].boxes[0].conf
+    # Print the evaluation results
+    print(results)
 
-# # Hook to store the activations
+if __name__ == "__main__":
-# activations = {}
+    parser = argparse.ArgumentParser(description='Train YOLOv10 model with PGT segmentation.')
+    parser.add_argument('--device', type=str, default='1', help='CUDA device number')
+    parser.add_argument('--batch_size', type=int, default=64, help='Batch size for training')
+    parser.add_argument('--epochs', type=int, default=100, help='Number of epochs for training')
+    parser.add_argument('--data_yaml', type=str, default='coco.yaml', help='Path to the data YAML file')
+    parser.add_argument('--model_path', type=str, default='yolov10n.pt', help='Path to the model file')
+    args = parser.parse_args()
 
-# def get_activation(name):
+    # Set CUDA device (only needed for multi-gpu machines)
-#     def hook(model, input, output):
+    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-#         activations[name] = output
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.device
-#     return hook
+    main(args)
-
-# # Register hooks for each layer you want to inspect
-# for name, layer in model.model.named_modules():
-#     layer.register_forward_hook(get_activation(name))
-
-# # Run the model to get activations
-# results = model.predict(image_tensor, save=True, visualize=True)
-
-# # # Print the activations
-# # for name, activation in activations.items():
-# #     print(f"Activation from layer {name}: {activation}")
-
-# # List activation names separately
-# print("\nActivation layer names:")
-# for name in activations.keys():
-#     print(name)
-# # pred.backward()
-
-# # Assuming 'model.23' is the layer of interest for bbox prediction and confidence
-# activation = activations['model.23']['one2one'][0]
-# act_23 = activations['model.23.cv3.2']
-# act_dfl = activations['model.23.dfl.conv']
-# act_conv = activations['model.0.conv']
-# act_act = activations['model.0.act']
-
-# # with torch.autograd.set_detect_anomaly(True):
-# #     pred.backward()
-# grad = torch.autograd.grad(act_23, im, grad_outputs=torch.ones_like(act_23), create_graph=True, retain_graph=True)[0]
-# # grad = torch.autograd.grad(pred, im, grad_outputs=torch.ones_like(pred), create_graph=True)[0]
-# grad = torch.autograd.grad(activations['model.23']['one2one'][1][0], 
-#                            activations['model.23.one2one_cv3.2'], 
-#                            grad_outputs=torch.ones_like(activations['model.23']['one2one'][1][0]), 
-#                            create_graph=True, retain_graph=True)[0]
-
-# # Print the results
-# print(results)
-
-# model.val(data='coco.yaml', batch=256)

ultralytics/models/yolo/segment/pgt_train.py

@@ -41,10 +41,8 @@ class PGTSegmentationTrainer(yolo.detect.PGTDetectionTrainer):
 
     def get_model(self, cfg=None, weights=None, verbose=True):
         """Return SegmentationModel initialized with specified config and weights."""
-        if self.args.model in ['yolov10n.pt', 'yolov10m.pt', 'yolov10x.pt', 'yolov10s.pt', 'yolov10b.pt', 'yolov10l.pt']:
+
         model = YOLOv10PGTDetectionModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
-        else:
-            model = SegmentationModel(cfg, ch=3, nc=self.data["nc"], verbose=verbose and RANK == -1)
         if weights:
             model.load(weights)
 
@@ -53,16 +51,10 @@ class PGTSegmentationTrainer(yolo.detect.PGTDetectionTrainer):
     def get_validator(self):
         """Return an instance of SegmentationValidator for validation of YOLO model."""
 
-        if self.args.model in ['yolov10n.pt', 'yolov10m.pt', 'yolov10x.pt', 'yolov10s.pt', 'yolov10b.pt', 'yolov10l.pt']:
         self.loss_names = "box_om", "cls_om", "dfl_om", "box_oo", "cls_oo", "dfl_oo", "pgt_loss",
         return YOLOv10PGTDetectionValidator(
             self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
-        else:
-            self.loss_names = "box_loss", "seg_loss", "cls_loss", "dfl_loss"
-            return yolo.segment.SegmentationValidator(
-                self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
-            )
 
     def plot_training_samples(self, batch, ni):
         """Creates a plot of training sample images with labels and box coordinates."""