diff --git a/run_val.py b/run_val.py index 98c5a476..40d5c499 100644 --- a/run_val.py +++ b/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) -os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" -os.environ["CUDA_VISIBLE_DEVICES"] = "4" +# nohup python run_pgt_train.py --device 1 > ./output_logs/gpu1_yolov10_pgt_train.log 2>&1 & -# model = YOLOv10() -# 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') +def main(args): -# Evaluate the model on the validation set -results = model.val(data='coco.yaml') + model = YOLOv10PGT(args.model_path) + + # Evaluate the model on the validation set + results = model.val(data=args.data_yaml) + + # Print the evaluation results + print(results) -# Print the evaluation results -print(results) +if __name__ == "__main__": + 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() -# pred = results[0].boxes[0].conf - -# # Hook to store the activations -# activations = {} - -# def get_activation(name): -# def hook(model, input, output): -# activations[name] = output -# return hook - -# # 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) \ No newline at end of file + # Set CUDA device (only needed for multi-gpu machines) + os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + os.environ["CUDA_VISIBLE_DEVICES"] = args.device + main(args) \ No newline at end of file diff --git a/ultralytics/models/yolo/segment/pgt_train.py b/ultralytics/models/yolo/segment/pgt_train.py index aa0559ee..6d3b4150 100644 --- a/ultralytics/models/yolo/segment/pgt_train.py +++ b/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) + + model = YOLOv10PGTDetectionModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1) if weights: model.load(weights) @@ -52,17 +50,11 @@ 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 - ) + + 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 + ) def plot_training_samples(self, batch, ni): """Creates a plot of training sample images with labels and box coordinates."""