def __init__(self, config, model_path): """ Initialization of PytorchYolo :param config: vision config dict :param model_path: path to the yolo model """ weightpath = os.path.join(model_path, "yolo_weights.weights") configpath = os.path.join(model_path, "config.cfg") self.model = torch_models.load_model(configpath, weightpath) self._image = None super(YoloHandlerPytorch, self).__init__(config, model_path)
def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_size=8, img_size=416, n_cpu=8, iou_thres=0.5, conf_thres=0.5, nms_thres=0.5, verbose=True): """Evaluate model on validation dataset. :param model_path: Path to model definition file (.cfg) :type model_path: str :param weights_path: Path to weights or checkpoint file (.weights or .pth) :type weights_path: str :param img_path: Path to file containing all paths to validation images. :type img_path: str :param class_names: List of class names :type class_names: [str] :param batch_size: Size of each image batch, defaults to 8 :type batch_size: int, optional :param img_size: Size of each image dimension for yolo, defaults to 416 :type img_size: int, optional :param n_cpu: Number of cpu threads to use during batch generation, defaults to 8 :type n_cpu: int, optional :param iou_thres: IOU threshold required to qualify as detected, defaults to 0.5 :type iou_thres: float, optional :param conf_thres: Object confidence threshold, defaults to 0.5 :type conf_thres: float, optional :param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5 :type nms_thres: float, optional :param verbose: If True, prints stats of model, defaults to True :type verbose: bool, optional :return: Returns precision, recall, AP, f1, ap_class """ dataloader = _create_validation_data_loader(img_path, batch_size, img_size, n_cpu) model = load_model(model_path, weights_path) metrics_output = _evaluate(model, dataloader, class_names, img_size, iou_thres, conf_thres, nms_thres, verbose) return metrics_output
def detect_directory(model_path, weights_path, img_path, classes, output_path, batch_size=8, img_size=416, n_cpu=8, conf_thres=0.5, nms_thres=0.5): """Detects objects on all images in specified directory and saves output images with drawn detections. :param model_path: Path to model definition file (.cfg) :type model_path: str :param weights_path: Path to weights or checkpoint file (.weights or .pth) :type weights_path: str :param img_path: Path to directory with images to inference :type img_path: str :param classes: List of class names :type classes: [str] :param output_path: Path to output directory :type output_path: str :param batch_size: Size of each image batch, defaults to 8 :type batch_size: int, optional :param img_size: Size of each image dimension for yolo, defaults to 416 :type img_size: int, optional :param n_cpu: Number of cpu threads to use during batch generation, defaults to 8 :type n_cpu: int, optional :param conf_thres: Object confidence threshold, defaults to 0.5 :type conf_thres: float, optional :param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5 :type nms_thres: float, optional """ dataloader = _create_data_loader(img_path, batch_size, img_size, n_cpu) model = load_model(model_path, weights_path) img_detections, imgs = detect(model, dataloader, output_path, img_size, conf_thres, nms_thres) _draw_and_save_output_images(img_detections, imgs, img_size, output_path, classes)
def run(): print_environment_info() parser = argparse.ArgumentParser(description="Trains the YOLO model.") parser.add_argument("-m", "--model", type=str, default="config/yolov3.cfg", help="Path to model definition file (.cfg)") parser.add_argument("-d", "--data", type=str, default="config/coco.data", help="Path to data config file (.data)") parser.add_argument("-e", "--epochs", type=int, default=300, help="Number of epochs") parser.add_argument("-v", "--verbose", action='store_true', help="Makes the training more verbose") parser.add_argument( "--n_cpu", type=int, default=8, help="Number of cpu threads to use during batch generation") parser.add_argument( "--pretrained_weights", type=str, help= "Path to checkpoint file (.weights or .pth). Starts training from checkpoint model" ) parser.add_argument("--checkpoint_interval", type=int, default=1, help="Interval of epochs between saving model weights") parser.add_argument( "--evaluation_interval", type=int, default=1, help="Interval of epochs between evaluations on validation set") parser.add_argument("--multiscale_training", action="store_false", help="Allow for multi-scale training") parser.add_argument( "--iou_thres", type=float, default=0.5, help="Evaluation: IOU threshold required to qualify as detected") parser.add_argument("--conf_thres", type=float, default=0.1, help="Evaluation: Object confidence threshold") parser.add_argument( "--nms_thres", type=float, default=0.5, help="Evaluation: IOU threshold for non-maximum suppression") parser.add_argument( "--logdir", type=str, default="logs", help="Directory for training log files (e.g. for TensorBoard)") parser.add_argument("--seed", type=int, default=-1, help="Makes results reproducable. Set -1 to disable.") args = parser.parse_args() print(f"Command line arguments: {args}") if args.seed != -1: provide_determinism(args.seed) logger = Logger(args.logdir) # Tensorboard logger # Create output directories if missing os.makedirs("output", exist_ok=True) os.makedirs("checkpoints", exist_ok=True) # Get data configuration data_config = parse_data_config(args.data) train_path = data_config["train"] valid_path = data_config["valid"] class_names = load_classes(data_config["names"]) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ############ # Create model # ############ model = load_model(args.model, args.pretrained_weights) # Print model if args.verbose: summary(model, input_size=(3, model.hyperparams['height'], model.hyperparams['height'])) mini_batch_size = model.hyperparams['batch'] // model.hyperparams[ 'subdivisions'] # ################# # Create Dataloader # ################# # Load training dataloader dataloader = _create_data_loader(train_path, mini_batch_size, model.hyperparams['height'], args.n_cpu, args.multiscale_training) # Load validation dataloader validation_dataloader = _create_validation_data_loader( valid_path, mini_batch_size, model.hyperparams['height'], args.n_cpu) # ################ # Create optimizer # ################ params = [p for p in model.parameters() if p.requires_grad] if (model.hyperparams['optimizer'] in [None, "adam"]): optimizer = optim.Adam( params, lr=model.hyperparams['learning_rate'], weight_decay=model.hyperparams['decay'], ) elif (model.hyperparams['optimizer'] == "sgd"): optimizer = optim.SGD(params, lr=model.hyperparams['learning_rate'], weight_decay=model.hyperparams['decay'], momentum=model.hyperparams['momentum']) else: print("Unknown optimizer. Please choose between (adam, sgd).") for epoch in range(args.epochs): print("\n---- Training Model ----") model.train() # Set model to training mode for batch_i, (_, imgs, targets) in enumerate( tqdm.tqdm(dataloader, desc=f"Training Epoch {epoch}")): batches_done = len(dataloader) * epoch + batch_i imgs = imgs.to(device, non_blocking=True) targets = targets.to(device) outputs = model(imgs) loss, loss_components = compute_loss(outputs, targets, model) loss.backward() ############### # Run optimizer ############### if batches_done % model.hyperparams['subdivisions'] == 0: # Adapt learning rate # Get learning rate defined in cfg lr = model.hyperparams['learning_rate'] if batches_done < model.hyperparams['burn_in']: # Burn in lr *= (batches_done / model.hyperparams['burn_in']) else: # Set and parse the learning rate to the steps defined in the cfg for threshold, value in model.hyperparams['lr_steps']: if batches_done > threshold: lr *= value # Log the learning rate logger.scalar_summary("train/learning_rate", lr, batches_done) # Set learning rate for g in optimizer.param_groups: g['lr'] = lr # Run optimizer optimizer.step() # Reset gradients optimizer.zero_grad() # ############ # Log progress # ############ if args.verbose: print( AsciiTable([ ["Type", "Value"], ["IoU loss", float(loss_components[0])], ["Object loss", float(loss_components[1])], ["Class loss", float(loss_components[2])], ["Loss", float(loss_components[3])], ["Batch loss", to_cpu(loss).item()], ]).table) # Tensorboard logging tensorboard_log = [("train/iou_loss", float(loss_components[0])), ("train/obj_loss", float(loss_components[1])), ("train/class_loss", float(loss_components[2])), ("train/loss", to_cpu(loss).item())] logger.list_of_scalars_summary(tensorboard_log, batches_done) model.seen += imgs.size(0) # ############# # Save progress # ############# # Save model to checkpoint file if epoch % args.checkpoint_interval == 0: checkpoint_path = f"checkpoints/yolov3_ckpt_{epoch}.pth" print(f"---- Saving checkpoint to: '{checkpoint_path}' ----") torch.save(model.state_dict(), checkpoint_path) # ######## # Evaluate # ######## if epoch % args.evaluation_interval == 0: print("\n---- Evaluating Model ----") # Evaluate the model on the validation set metrics_output = _evaluate(model, validation_dataloader, class_names, img_size=model.hyperparams['height'], iou_thres=args.iou_thres, conf_thres=args.conf_thres, nms_thres=args.nms_thres, verbose=args.verbose) if metrics_output is not None: precision, recall, AP, f1, ap_class = metrics_output evaluation_metrics = [("validation/precision", precision.mean()), ("validation/recall", recall.mean()), ("validation/mAP", AP.mean()), ("validation/f1", f1.mean())] logger.list_of_scalars_summary(evaluation_metrics, epoch)
import torch from pytorchyolo import models from nni.compression.pytorch import ModelSpeedup from nni.algorithms.compression.pytorch.pruning import L1FilterPruner, LevelPruner from nni.compression.pytorch.utils import not_safe_to_prune # The Yolo can be downloaded at https://github.com/eriklindernoren/PyTorch-YOLOv3.git prefix = '/home/user/PyTorch-YOLOv3' # replace this path with yours # Load the YOLO model model = models.load_model( "%s/config/yolov3.cfg" % prefix, "%s/yolov3.weights" % prefix).cpu() model.eval() dummy_input = torch.rand(8, 3, 320, 320) model(dummy_input) # Generate the config list for pruner # Filter the layers that may not be able to prune not_safe = not_safe_to_prune(model, dummy_input) cfg_list = [] for name, module in model.named_modules(): if name in not_safe: continue if isinstance(module, torch.nn.Conv2d): cfg_list.append({'op_types':['Conv2d'], 'sparsity':0.6, 'op_names':[name]}) # Prune the model pruner = L1FilterPruner(model, cfg_list) pruner.compress() pruner.export_model('./model', './mask') pruner._unwrap_model() # Speedup the model
from pytorchyolo.utils.loss import compute_loss from pytorchyolo.models import load_model from pytorchyolo.models import Darknet model = Net() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) learning_rate = 1e-6 params = [p for p in model.parameters() if p.requires_grad] optimizer = optim.Adam(lr=model.hyperparams['learning rate'], params=params) epochs = 10 for epoch in range(epochs): print('training model....') model = load_model(model_path='yolo_model.py') model.train() for batch_i, (_, imgs, targets) in enumerate(tqdm.tqdm(dataloader, desc=f"Training Epoch{epoch}")): batches_done = len(dataloader) * epoch + batch_i imgs = imgs.to(device, non_blocking=True) targets = targets.to(device) outputs = model(imgs) loss, loss_components = compute_loss(outputs, targets, model) loss.backward() optimizer.step() optimizer.zero_grad() print(epoch+1)